prefecthq--fastmcp/tests/server/sampling/test_sampling_tool.py

"""Tests for SamplingTool."""

import pytest
from mcp.server.auth.middleware.auth_context import auth_context_var
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser

from fastmcp.exceptions import AuthorizationError
from fastmcp.server.auth import AccessToken, require_scopes
from fastmcp.server.context import _current_transport
from fastmcp.server.sampling import SamplingTool
from fastmcp.tools.function_tool import FunctionTool
from fastmcp.tools.tool_transform import ArgTransform, TransformedTool


class TestSamplingToolFromFunction:
    """Tests for SamplingTool.from_function()."""

    def test_from_simple_function(self):
        def search(query: str) -> str:
            """Search the web."""
            return f"Results for: {query}"

        tool = SamplingTool.from_function(search)

        assert tool.name == "search"
        assert tool.description == "Search the web."
        assert "query" in tool.parameters.get("properties", {})
        assert tool.fn is search

    def test_from_function_with_overrides(self):
        def search(query: str) -> str:
            return f"Results for: {query}"

        tool = SamplingTool.from_function(
            search,
            name="web_search",
            description="Search the internet",
        )

        assert tool.name == "web_search"
        assert tool.description == "Search the internet"

    def test_from_lambda_requires_name(self):
        with pytest.raises(ValueError, match="must provide a name for lambda"):
            SamplingTool.from_function(lambda x: x)

    def test_from_lambda_with_name(self):
        tool = SamplingTool.from_function(lambda x: x * 2, name="double")

        assert tool.name == "double"

    def test_from_async_function(self):
        async def async_search(query: str) -> str:
            """Async search."""
            return f"Async results for: {query}"

        tool = SamplingTool.from_function(async_search)

        assert tool.name == "async_search"
        assert tool.description == "Async search."

    def test_multiple_parameters(self):
        def search(query: str, limit: int = 10, include_images: bool = False) -> str:
            """Search with options."""
            return f"Results for: {query}"

        tool = SamplingTool.from_function(search)
        props = tool.parameters.get("properties", {})

        assert "query" in props
        assert "limit" in props
        assert "include_images" in props


class TestSamplingToolRun:
    """Tests for SamplingTool.run()."""

    async def test_run_sync_function(self):
        def add(a: int, b: int) -> int:
            """Add two numbers."""
            return a + b

        tool = SamplingTool.from_function(add)
        result = await tool.run({"a": 2, "b": 3})
        assert result == 5

    async def test_run_async_function(self):
        async def async_add(a: int, b: int) -> int:
            """Add two numbers asynchronously."""
            return a + b

        tool = SamplingTool.from_function(async_add)
        result = await tool.run({"a": 2, "b": 3})
        assert result == 5

    async def test_run_with_no_arguments(self):
        def get_value() -> str:
            """Return a fixed value."""
            return "hello"

        tool = SamplingTool.from_function(get_value)
        result = await tool.run()
        assert result == "hello"

    async def test_run_with_none_arguments(self):
        def get_value() -> str:
            """Return a fixed value."""
            return "hello"

        tool = SamplingTool.from_function(get_value)
        result = await tool.run(None)
        assert result == "hello"


class TestSamplingToolSDKConversion:
    """Tests for SamplingTool._to_sdk_tool() internal method."""

    def test_to_sdk_tool(self):
        def search(query: str) -> str:
            """Search the web."""
            return f"Results for: {query}"

        tool = SamplingTool.from_function(search)
        sdk_tool = tool._to_sdk_tool()

        assert sdk_tool.name == "search"
        assert sdk_tool.description == "Search the web."
        assert "query" in sdk_tool.input_schema.get("properties", {})


class TestSamplingToolFromCallableTool:
    """Tests for SamplingTool.from_callable_tool()."""

    def test_from_function_tool(self):
        """Test converting a FunctionTool to SamplingTool."""

        def search(query: str) -> str:
            """Search the web."""
            return f"Results for: {query}"

        function_tool = FunctionTool.from_function(search)
        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        assert sampling_tool.name == "search"
        assert sampling_tool.description == "Search the web."
        assert "query" in sampling_tool.parameters.get("properties", {})
        # fn is now a wrapper that calls tool.run() for proper result processing
        assert callable(sampling_tool.fn)

    def test_from_function_tool_with_overrides(self):
        """Test converting FunctionTool with name/description overrides."""

        def search(query: str) -> str:
            """Search the web."""
            return f"Results for: {query}"

        function_tool = FunctionTool.from_function(search)
        sampling_tool = SamplingTool.from_callable_tool(
            function_tool,
            name="web_search",
            description="Search the internet",
        )

        assert sampling_tool.name == "web_search"
        assert sampling_tool.description == "Search the internet"

    def test_from_transformed_tool(self):
        """Test converting a TransformedTool to SamplingTool."""

        def original(query: str, limit: int) -> str:
            """Original tool."""
            return f"Results for: {query} (limit: {limit})"

        function_tool = FunctionTool.from_function(original)
        transformed_tool = TransformedTool.from_tool(
            function_tool,
            name="search_transformed",
            transform_args={"query": ArgTransform(name="q")},
        )

        sampling_tool = SamplingTool.from_callable_tool(transformed_tool)

        assert sampling_tool.name == "search_transformed"
        assert sampling_tool.description == "Original tool."
        # The transformed tool should have 'q' instead of 'query'
        assert "q" in sampling_tool.parameters.get("properties", {})
        assert "limit" in sampling_tool.parameters.get("properties", {})

    async def test_from_function_tool_execution(self):
        """Test that converted FunctionTool executes correctly."""

        def add(a: int, b: int) -> int:
            """Add two numbers."""
            return a + b

        function_tool = FunctionTool.from_function(add)
        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        result = await sampling_tool.run({"a": 2, "b": 3})
        assert result == 5

    async def test_from_transformed_tool_execution(self):
        """Test that converted TransformedTool executes correctly."""

        def multiply(x: int, y: int) -> int:
            """Multiply two numbers."""
            return x * y

        function_tool = FunctionTool.from_function(multiply)
        transformed_tool = TransformedTool.from_tool(
            function_tool,
            transform_args={"x": ArgTransform(name="a"), "y": ArgTransform(name="b")},
        )

        sampling_tool = SamplingTool.from_callable_tool(transformed_tool)

        # Use the transformed parameter names
        result = await sampling_tool.run({"a": 3, "b": 4})
        # Result should be unwrapped from ToolResult
        assert result == 12

    def test_from_invalid_tool_type(self):
        """Test that from_callable_tool rejects non-tool objects."""

        class NotATool:
            pass

        with pytest.raises(
            TypeError,
            match="Expected FunctionTool or TransformedTool",
        ):
            SamplingTool.from_callable_tool(NotATool())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]

    def test_from_plain_function_fails(self):
        """Test that plain functions are rejected by from_callable_tool."""

        def my_function():
            pass

        with pytest.raises(TypeError, match="Expected FunctionTool or TransformedTool"):
            SamplingTool.from_callable_tool(my_function)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]

    async def test_from_function_tool_with_output_schema(self):
        """Test that FunctionTool with output_schema is handled correctly."""

        def search(query: str) -> dict:
            """Search for something."""
            return {"results": ["item1", "item2"], "count": 2}

        # Create FunctionTool with x-fastmcp-wrap-result
        function_tool = FunctionTool.from_function(
            search,
            output_schema={
                "type": "object",
                "properties": {
                    "results": {"type": "array"},
                    "count": {"type": "integer"},
                },
                "x-fastmcp-wrap-result": True,
            },
        )

        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        # Run the tool - should unwrap the {"result": {...}} wrapper
        result = await sampling_tool.run({"query": "test"})

        # Should get the unwrapped dict, not ToolResult
        assert isinstance(result, dict)
        assert result == {"results": ["item1", "item2"], "count": 2}

    async def test_from_function_tool_without_wrap_result(self):
        """Test that FunctionTool without x-fastmcp-wrap-result is handled correctly."""

        def get_data() -> dict:
            """Get some data."""
            return {"status": "ok", "value": 42}

        # Create FunctionTool with output_schema but no wrap-result flag
        function_tool = FunctionTool.from_function(
            get_data,
            output_schema={
                "type": "object",
                "properties": {
                    "status": {"type": "string"},
                    "value": {"type": "integer"},
                },
            },
        )

        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        # Run the tool - should return structured_content directly
        result = await sampling_tool.run({})

        assert isinstance(result, dict)
        assert result == {"status": "ok", "value": 42}


class TestSamplingToolAuthEnforcement:
    """Tests that auth-protected tools enforce auth when used via sampling."""

    async def test_auth_protected_tool_blocked_without_token(self):
        """An auth-protected tool wrapped as SamplingTool must reject
        calls when no valid token is present in a non-stdio transport."""

        def secret_action() -> str:
            """Do something privileged."""
            return "secret"

        function_tool = FunctionTool.from_function(
            secret_action,
            auth=require_scopes("admin"),
        )
        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        transport_token = _current_transport.set("streamable-http")
        try:
            with pytest.raises(AuthorizationError, match="insufficient permissions"):
                await sampling_tool.run({})
        finally:
            _current_transport.reset(transport_token)

    async def test_auth_protected_tool_blocked_with_wrong_scopes(self):
        """An auth-protected tool rejects calls when the token lacks
        the required scopes."""

        def secret_action() -> str:
            """Do something privileged."""
            return "secret"

        function_tool = FunctionTool.from_function(
            secret_action,
            auth=require_scopes("admin"),
        )
        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        token = AccessToken(
            token="test",
            client_id="c",
            scopes=["read"],
            expires_at=None,
            claims={},
        )
        transport_token = _current_transport.set("streamable-http")
        auth_token = auth_context_var.set(AuthenticatedUser(token))
        try:
            with pytest.raises(AuthorizationError, match="insufficient permissions"):
                await sampling_tool.run({})
        finally:
            auth_context_var.reset(auth_token)
            _current_transport.reset(transport_token)

    async def test_auth_protected_tool_allowed_with_correct_scopes(self):
        """An auth-protected tool succeeds when the token has the
        required scopes."""

        def secret_action() -> str:
            """Do something privileged."""
            return "secret"

        function_tool = FunctionTool.from_function(
            secret_action,
            auth=require_scopes("admin"),
        )
        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        token = AccessToken(
            token="test",
            client_id="c",
            scopes=["admin"],
            expires_at=None,
            claims={},
        )
        transport_token = _current_transport.set("streamable-http")
        auth_token = auth_context_var.set(AuthenticatedUser(token))
        try:
            result = await sampling_tool.run({})
            assert result == "secret"
        finally:
            auth_context_var.reset(auth_token)
            _current_transport.reset(transport_token)

    async def test_auth_protected_tool_skipped_on_stdio(self):
        """Auth checks are skipped for stdio transport, matching
        server dispatcher behavior."""

        def secret_action() -> str:
            """Do something privileged."""
            return "secret"

        function_tool = FunctionTool.from_function(
            secret_action,
            auth=require_scopes("admin"),
        )
        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        transport_token = _current_transport.set("stdio")
        try:
            result = await sampling_tool.run({})
            assert result == "secret"
        finally:
            _current_transport.reset(transport_token)

    async def test_tool_without_auth_runs_normally(self):
        """Tools without auth still run without any auth context."""

        def public_action() -> str:
            """Do something public."""
            return "public"

        function_tool = FunctionTool.from_function(public_action)
        sampling_tool = SamplingTool.from_callable_tool(function_tool)

        result = await sampling_tool.run({})
        assert result == "public"

    async def test_auth_protected_transformed_tool_blocked(self):
        """Auth checks also apply to TransformedTools with auth."""

        def secret_action(x: int) -> int:
            """Privileged computation."""
            return x * 2

        function_tool = FunctionTool.from_function(
            secret_action,
            auth=require_scopes("compute"),
        )
        transformed_tool = TransformedTool.from_tool(
            function_tool,
            transform_args={"x": ArgTransform(name="value")},
        )
        sampling_tool = SamplingTool.from_callable_tool(transformed_tool)

        transport_token = _current_transport.set("streamable-http")
        try:
            with pytest.raises(AuthorizationError, match="insufficient permissions"):
                await sampling_tool.run({"value": 5})
        finally:
            _current_transport.reset(transport_token)