Proxy support advanced MCP features (#1022)

* Proxy support advanced MCP features

* Optimize ProxyClient testing

* Add documentation for ProxyClient

tests/server/proxy/test_proxy_client.py

@@ -0,0 +1,257 @@
+from dataclasses import dataclass
+from typing import cast
+
+import pytest
+from mcp.types import LoggingLevel, ModelHint, ModelPreferences, TextContent
+
+from fastmcp import Client, Context, FastMCP
+from fastmcp.client.elicitation import ElicitRequestParams, ElicitResult
+from fastmcp.client.logging import LogMessage
+from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
+from fastmcp.exceptions import ToolError
+from fastmcp.server.proxy import ProxyClient
+
+
+@pytest.fixture
+def fastmcp_server():
+    mcp = FastMCP("TestServer")
+
+    @mcp.tool
+    async def list_roots(context: Context) -> list[str]:
+        roots = await context.list_roots()
+        return [str(r.uri) for r in roots]
+
+    @mcp.tool
+    async def sampling(
+        context: Context,
+    ) -> str:
+        result = await context.sample(
+            "Hello, world!",
+            system_prompt="You love FastMCP",
+            temperature=0.5,
+            max_tokens=100,
+            model_preferences="gpt-4o",
+        )
+        return cast(TextContent, result).text
+
+    @dataclass
+    class Person:
+        name: str
+
+    @mcp.tool
+    async def elicit(context: Context) -> str:
+        result = await context.elicit(
+            message="What is your name?",
+            response_type=Person,
+        )
+        if result.action == "accept":
+            return f"Hello, {result.data.name}!"
+        else:
+            return "No name provided."
+
+    @mcp.tool
+    async def log(
+        message: str, level: LoggingLevel, logger: str, context: Context
+    ) -> None:
+        await context.log(message=message, level=level, logger_name=logger)
+
+    @mcp.tool
+    async def report_progress(context: Context) -> int:
+        for i in range(3):
+            await context.report_progress(
+                progress=i + 1,
+                total=3,
+                message=f"{(i + 1) / 3 * 100:.2f}% complete",
+            )
+        return 100
+
+    return mcp
+
+
+@pytest.fixture
+async def proxy_server(fastmcp_server: FastMCP):
+    """
+    A proxy server that forwards interactions with the proxy client to the given fastmcp server.
+    """
+    return FastMCP.as_proxy(ProxyClient(fastmcp_server))
+
+
+class TestProxyClient:
+    async def test_forward_error_response(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards an error response.
+        """
+        async with Client(proxy_server) as client:
+            with pytest.raises(ToolError, match="Elicitation not supported"):
+                await client.call_tool("elicit", {})
+
+    async def test_forward_list_roots_request(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `list_roots` request.
+        """
+        roots_handler_called = False
+
+        async def roots_handler(ctx: RequestContext):
+            nonlocal roots_handler_called
+            roots_handler_called = True
+            return []
+
+        async with Client(proxy_server, roots=roots_handler) as client:
+            await client.call_tool("list_roots", {})
+
+        assert roots_handler_called
+
+    async def test_forward_list_roots_response(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `list_roots` response.
+        """
+        async with Client(proxy_server, roots=["file://x/y/z"]) as client:
+            result = await client.call_tool("list_roots", {})
+            assert result.data == ["file://x/y/z"]
+
+    async def test_forward_sampling_request(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `sampling` request.
+        """
+        sampling_handler_called = False
+
+        def sampling_handler(
+            messages: list[SamplingMessage],
+            params: SamplingParams,
+            ctx: RequestContext,
+        ) -> str:
+            nonlocal sampling_handler_called
+            sampling_handler_called = True
+            assert messages == [
+                SamplingMessage(
+                    role="user",
+                    content=TextContent(type="text", text="Hello, world!"),
+                )
+            ]
+            assert params.systemPrompt == "You love FastMCP"
+            assert params.temperature == 0.5
+            assert params.maxTokens == 100
+            assert params.modelPreferences == ModelPreferences(
+                hints=[ModelHint(name="gpt-4o")]
+            )
+            return ""
+
+        async with Client(proxy_server, sampling_handler=sampling_handler) as client:
+            await client.call_tool("sampling", {})
+
+        assert sampling_handler_called
+
+    async def test_forward_sampling_response(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `sampling` response.
+        """
+        async with Client(
+            proxy_server, sampling_handler=lambda *args: "I love FastMCP"
+        ) as client:
+            result = await client.call_tool("sampling", {})
+            assert result.data == "I love FastMCP"
+
+    async def test_elicit_request(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `elicit` request.
+        """
+        elicitation_handler_called = False
+
+        async def elicitation_handler(
+            message, response_type, params: ElicitRequestParams, ctx
+        ):
+            nonlocal elicitation_handler_called
+            elicitation_handler_called = True
+            assert message == "What is your name?"
+            assert "Person" in str(response_type)
+            assert params.requestedSchema == {
+                "title": "Person",
+                "type": "object",
+                "properties": {"name": {"title": "Name", "type": "string"}},
+                "required": ["name"],
+            }
+            return ElicitResult(action="accept", content=response_type(name="Alice"))
+
+        async with Client(
+            proxy_server, elicitation_handler=elicitation_handler
+        ) as client:
+            await client.call_tool("elicit", {})
+
+        assert elicitation_handler_called
+
+    async def test_elicit_accept_response(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `elicit` accept response.
+        """
+
+        async def elicitation_handler(
+            message, response_type, params: ElicitRequestParams, ctx
+        ):
+            return ElicitResult(action="accept", content=response_type(name="Alice"))
+
+        async with Client(
+            proxy_server,
+            elicitation_handler=elicitation_handler,
+        ) as client:
+            result = await client.call_tool("elicit", {})
+            assert result.data == "Hello, Alice!"
+
+    async def test_elicit_decline_response(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `elicit` decline response.
+        """
+
+        async def elicitation_handler(
+            message, response_type, params: ElicitRequestParams, ctx
+        ):
+            return ElicitResult(action="decline")
+
+        async with Client(
+            proxy_server, elicitation_handler=elicitation_handler
+        ) as client:
+            result = await client.call_tool("elicit", {})
+            assert result.data == "No name provided."
+
+    async def test_log_request(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `log` request.
+        """
+        log_handler_called = False
+
+        async def log_handler(message: LogMessage) -> None:
+            nonlocal log_handler_called
+            log_handler_called = True
+            assert message.data == "Hello, world!"
+            assert message.level == "info"
+            assert message.logger == "test"
+
+        async with Client(proxy_server, log_handler=log_handler) as client:
+            await client.call_tool(
+                "log", {"message": "Hello, world!", "level": "info", "logger": "test"}
+            )
+
+        assert log_handler_called
+
+    async def test_report_progress_request(self, proxy_server: FastMCP):
+        """
+        Test that the proxy client correctly forwards the `report_progress` request.
+        """
+
+        EXPECTED_PROGRESS_MESSAGES = [
+            dict(progress=1, total=3, message="33.33% complete"),
+            dict(progress=2, total=3, message="66.67% complete"),
+            dict(progress=3, total=3, message="100.00% complete"),
+        ]
+        PROGRESS_MESSAGES = []
+
+        async def progress_handler(
+            progress: float, total: float | None, message: str | None
+        ) -> None:
+            PROGRESS_MESSAGES.append(
+                dict(progress=progress, total=total, message=message)
+            )
+
+        async with Client(proxy_server, progress_handler=progress_handler) as client:
+            await client.call_tool("report_progress", {})
+
+        assert PROGRESS_MESSAGES == EXPECTED_PROGRESS_MESSAGES