prefecthq--fastmcp/tests/server/providers/proxy/test_proxy_server.py

import inspect
import json
import time
from typing import Any, cast
from unittest.mock import AsyncMock, patch

import mcp_types
import pytest
from anyio import create_task_group
from dirty_equals import Contains
from mcp import MCPError
from mcp_types import Icon, TextContent, TextResourceContents
from pydantic import AnyUrl

from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.client.transports import FastMCPTransport, StreamableHttpTransport
from fastmcp.exceptions import ToolError
from fastmcp.resources import ResourceContent, ResourceResult
from fastmcp.server import create_proxy
from fastmcp.server.providers.proxy import (
    FastMCPProxy,
    ProxyClient,
    ProxyProvider,
)
from fastmcp.tools.base import ToolResult
from fastmcp.tools.tool_transform import (
    ToolTransformConfig,
)
from fastmcp.utilities.http import find_available_port
from fastmcp.utilities.tests import run_server_async

USERS = [
    {"id": "1", "name": "Alice", "active": True},
    {"id": "2", "name": "Bob", "active": True},
    {"id": "3", "name": "Charlie", "active": False},
]


@pytest.fixture
def fastmcp_server():
    server = FastMCP("TestServer")

    # --- Tools ---

    @server.tool(
        tags={"greet"},
        title="Greet",
        icons=[Icon(src="https://example.com/greet-icon.png")],
    )
    def greet(name: str) -> str:
        """Greet someone by name."""
        return f"Hello, {name}!"

    @server.tool
    def tool_without_description() -> str:
        return "Hello?"

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

    @server.tool
    def error_tool():
        """This tool always raises an error."""
        raise ValueError("This is a test error")

    # --- Resources ---

    @server.resource(
        uri="resource://wave",
        tags={"wave"},
        title="Wave",
        icons=[Icon(src="https://example.com/wave-icon.png")],
    )
    def wave() -> str:
        return "👋"

    @server.resource(uri="data://users")
    async def get_users() -> str:
        import json

        return json.dumps(USERS, separators=(",", ":"))

    @server.resource(
        uri="data://user/{user_id}",
        tags={"users"},
        title="User Template",
        icons=[Icon(src="https://example.com/user-icon.png")],
    )
    async def get_user(user_id: str) -> str:
        import json

        user = next((user for user in USERS if user["id"] == user_id), None)
        return json.dumps(user, separators=(",", ":")) if user else "null"

    @server.resource(uri="data://multi")
    def get_multi_content() -> ResourceResult:
        """Resource that returns multiple content items."""
        return ResourceResult(
            contents=[
                ResourceContent(content="First item", mime_type="text/plain"),
                ResourceContent(
                    content='{"key": "value"}', mime_type="application/json"
                ),
                ResourceContent(
                    content="# Markdown\nContent", mime_type="text/markdown"
                ),
            ],
            meta={"count": 3},
        )

    @server.resource(uri="data://multi/{id}")
    def get_multi_template(id: str) -> ResourceResult:
        """Resource template that returns multiple content items."""
        return ResourceResult(
            contents=[
                ResourceContent(content=f"Item {id} - First", mime_type="text/plain"),
                ResourceContent(
                    content=f'{{"id": "{id}", "status": "active"}}',
                    mime_type="application/json",
                ),
            ],
            meta={"id": id},
        )

    # --- Prompts ---

    @server.prompt(
        tags={"welcome"},
        title="Welcome",
        icons=[Icon(src="https://example.com/welcome-icon.png")],
    )
    def welcome(name: str) -> str:
        return f"Welcome to FastMCP, {name}!"

    @server.prompt
    def image_prompt():
        """A prompt that returns an image."""
        from fastmcp.prompts.base import Message, PromptResult

        return PromptResult(
            messages=[
                Message("Here is an image:"),
                Message(
                    content=mcp_types.ImageContent(
                        type="image",
                        data="iVBORw0KGgoAAAANSUhEUg==",
                        mime_type="image/png",
                    ),
                    role="user",
                ),
            ]
        )

    return server


@pytest.fixture
async def proxy_server(fastmcp_server):
    """Fixture that creates a FastMCP proxy server."""
    return create_proxy(ProxyClient(transport=FastMCPTransport(fastmcp_server)))


async def test_create_proxy_with_client(fastmcp_server):
    """Test create_proxy with a Client."""
    client = ProxyClient(transport=FastMCPTransport(fastmcp_server))
    server = create_proxy(client)

    assert isinstance(server, FastMCPProxy)
    assert isinstance(server, FastMCP)
    assert server.name.startswith("FastMCPProxy-")


async def test_create_proxy_with_server(fastmcp_server):
    """create_proxy should accept a FastMCP instance."""
    proxy = create_proxy(fastmcp_server)
    async with Client(proxy) as client:
        result = await client.call_tool("greet", {"name": "Test"})
        assert result.data == "Hello, Test!"


async def test_create_proxy_with_transport(fastmcp_server):
    """create_proxy should accept a ClientTransport."""
    proxy = create_proxy(FastMCPTransport(fastmcp_server))
    async with Client(proxy) as client:
        result = await client.call_tool("greet", {"name": "Test"})
        assert result.data == "Hello, Test!"


def test_create_proxy_with_url():
    """create_proxy should accept a URL without connecting."""
    proxy = create_proxy("http://example.com/mcp/")
    assert isinstance(proxy, FastMCPProxy)
    client = cast(Client, proxy.client_factory())
    assert isinstance(client.transport, StreamableHttpTransport)
    assert client.transport.url == "http://example.com/mcp/"


async def test_proxy_with_async_client_factory():
    """FastMCPProxy should accept an async client_factory."""

    async def async_factory():
        return Client("http://example.com/mcp/")

    proxy = FastMCPProxy(client_factory=async_factory)
    assert isinstance(proxy, FastMCPProxy)
    assert inspect.iscoroutinefunction(proxy.client_factory)
    client = proxy.client_factory()
    if inspect.isawaitable(client):
        client = await client
    assert isinstance(client, Client)
    assert isinstance(client.transport, StreamableHttpTransport)
    assert client.transport.url == "http://example.com/mcp/"


async def test_proxy_ping_forwards_to_remote_server(fastmcp_server):
    proxy = create_proxy(fastmcp_server)

    async with Client(proxy) as client:
        assert await client.ping() is True


async def test_proxy_ping_surfaces_wrong_remote_path():
    remote = FastMCP("remote")
    async with run_server_async(remote, transport="http") as url:
        proxy = create_proxy(StreamableHttpTransport(url.removesuffix("/mcp")))

        # SDK v2 surfaces a wrong remote path as an HTTP "Not Found" rather than
        # the v1 "Session terminated" message.
        with pytest.raises(MCPError, match="Not Found"):
            async with Client(proxy):
                pass


async def test_proxy_initialize_forwards_remote_connection_error():
    port = find_available_port()
    proxy = create_proxy(
        StreamableHttpTransport(f"http://127.0.0.1:{port}/mcp"),
        provider_error_strategy="raise",
    )

    with pytest.raises(MCPError, match="Client failed to connect"):
        async with Client(proxy):
            pass


async def test_proxy_list_tools_surfaces_remote_connection_error():
    port = find_available_port()
    proxy = create_proxy(
        StreamableHttpTransport(f"http://127.0.0.1:{port}/mcp"),
        provider_error_strategy="raise",
    )

    with pytest.raises(RuntimeError, match="Client failed to connect"):
        await proxy.list_tools()


async def test_proxy_list_tools_client_surfaces_remote_connection_error():
    port = find_available_port()
    proxy = create_proxy(
        StreamableHttpTransport(f"http://127.0.0.1:{port}/mcp"),
        provider_error_strategy="raise",
    )

    with pytest.raises(MCPError, match="Client failed to connect"):
        async with Client(proxy) as client:
            await client.list_tools()


class TestTools:
    async def test_get_tools(self, proxy_server):
        tools = await proxy_server.list_tools()
        assert any(t.name == "greet" for t in tools)
        assert any(t.name == "add" for t in tools)
        assert any(t.name == "error_tool" for t in tools)
        assert any(t.name == "tool_without_description" for t in tools)

    async def test_get_tools_meta(self, proxy_server):
        tools = await proxy_server.list_tools()
        greet_tool = next(t for t in tools if t.name == "greet")
        assert greet_tool.title == "Greet"
        assert greet_tool.meta == {"fastmcp": {"tags": ["greet"]}}
        assert greet_tool.icons == [Icon(src="https://example.com/greet-icon.png")]

    async def test_get_transformed_tools(self):
        """Test that tool transformations are applied to proxied tools."""
        from fastmcp.server.transforms import ToolTransform

        # Create server with transformation
        server = FastMCP("TestServer")

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

        server.add_transform(
            ToolTransform({"add": ToolTransformConfig(name="add_transformed")})
        )

        proxy = create_proxy(server)
        tools = await proxy.list_tools()
        assert any(t.name == "add_transformed" for t in tools)
        assert not any(t.name == "add" for t in tools)

    async def test_call_transformed_tools(self):
        """Test calling a transformed tool through a proxy."""
        from fastmcp.server.transforms import ToolTransform

        # Create server with transformation
        server = FastMCP("TestServer")

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

        server.add_transform(
            ToolTransform({"add": ToolTransformConfig(name="add_transformed")})
        )

        proxy = create_proxy(server)
        async with Client(proxy) as client:
            result = await client.call_tool("add_transformed", {"a": 1, "b": 2})
        assert result.data == 3

    async def test_tool_without_description(self, proxy_server):
        tools = await proxy_server.list_tools()
        tool = next(t for t in tools if t.name == "tool_without_description")
        assert tool.description is None

    async def test_list_tools_same_as_original(self, fastmcp_server, proxy_server):
        async with Client(fastmcp_server) as original_client:
            original = await original_client.list_tools()
        async with Client(proxy_server) as proxy_client:
            proxied = await proxy_client.list_tools()
        assert proxied == original

    async def test_call_tool_result_same_as_original(
        self, fastmcp_server: FastMCP, proxy_server: FastMCPProxy
    ):
        async with Client(fastmcp_server) as original_client:
            result = await original_client.call_tool("greet", {"name": "Alice"})
        async with Client(proxy_server) as proxy_client:
            proxy_result = await proxy_client.call_tool("greet", {"name": "Alice"})

        assert result.content == proxy_result.content
        assert result.data == proxy_result.data

    async def test_call_tool_calls_tool(self, proxy_server):
        async with Client(proxy_server) as client:
            proxy_result = await client.call_tool("add", {"a": 1, "b": 2})
        assert proxy_result.data == 3

    async def test_error_tool_raises_error(self, proxy_server):
        with pytest.raises(ToolError, match="This is a test error"):
            async with Client(proxy_server) as client:
                await client.call_tool("error_tool", {})

    async def test_error_tool_with_image_content(self, proxy_server):
        """Non-TextContent error responses should not crash with AttributeError."""
        error_result = mcp_types.CallToolResult(
            content=[
                mcp_types.ImageContent(
                    type="image", data="abc123", mime_type="image/png"
                )
            ],
            is_error=True,
        )
        with patch.object(
            Client, "call_tool_mcp", new_callable=AsyncMock, return_value=error_result
        ):
            with pytest.raises(ToolError):
                async with Client(proxy_server) as client:
                    await client.call_tool("error_tool", {})

    async def test_error_tool_with_empty_content(self, proxy_server):
        """Error responses with empty content should not crash."""
        error_result = mcp_types.CallToolResult(
            content=[],
            is_error=True,
        )
        with patch.object(
            Client, "call_tool_mcp", new_callable=AsyncMock, return_value=error_result
        ):
            with pytest.raises(ToolError):
                async with Client(proxy_server) as client:
                    await client.call_tool("error_tool", {})

    async def test_error_tool_passthrough_preserves_content(self, proxy_server):
        """Upstream error results pass through with content intact, not flattened."""
        error_result = mcp_types.CallToolResult(
            content=[
                mcp_types.ImageContent(
                    type="image", data="abc123", mime_type="image/png"
                )
            ],
            structured_content={"detail": "boom"},
            is_error=True,
        )
        with patch.object(
            Client, "call_tool_mcp", new_callable=AsyncMock, return_value=error_result
        ):
            async with Client(proxy_server) as client:
                result = await client.call_tool("error_tool", {}, raise_on_error=False)

        assert result.is_error is True
        assert isinstance(result.content[0], mcp_types.ImageContent)
        assert result.content[0].data == "abc123"
        assert result.structured_content == {"detail": "boom"}

    async def test_call_tool_forwards_meta(self, fastmcp_server, proxy_server):
        """Test that metadata from proxied tool results is properly forwarded."""

        @fastmcp_server.tool
        def tool_with_meta(value: str) -> ToolResult:
            """A tool that returns metadata in its result."""
            return ToolResult(
                content=f"Result: {value}",
                meta={"custom_key": "custom_value", "processed": True},
            )

        async with Client(proxy_server) as client:
            result = await client.call_tool("tool_with_meta", {"value": "test"})

        assert isinstance(result.content[0], TextContent)
        assert result.content[0].text == "Result: test"
        assert result.meta == {"custom_key": "custom_value", "processed": True}

    async def test_proxy_can_overwrite_proxied_tool(self, proxy_server):
        """
        Test that a tool defined on the proxy can overwrite the proxied tool with the same name.
        """

        @proxy_server.tool
        def greet(name: str, extra: str = "extra") -> str:
            return f"Overwritten, {name}! {extra}"

        async with Client(proxy_server) as client:
            result = await client.call_tool("greet", {"name": "Marvin", "extra": "abc"})
        assert result.data == "Overwritten, Marvin! abc"

    async def test_proxy_can_list_overwritten_tool(self, proxy_server):
        """
        Test that a tool defined on the proxy is listed instead of the proxied tool
        """

        @proxy_server.tool
        def greet(name: str, extra: str = "extra") -> str:
            return f"Overwritten, {name}! {extra}"

        async with Client(proxy_server) as client:
            tools = await client.list_tools()
            greet_tool = next(t for t in tools if t.name == "greet")
            assert "extra" in greet_tool.input_schema["properties"]


class TestResources:
    async def test_get_resources(self, proxy_server):
        resources = await proxy_server.list_resources()
        assert [r.uri for r in resources] == Contains(
            AnyUrl("data://users"),
            AnyUrl("resource://wave"),
        )
        assert [r.name for r in resources] == Contains("get_users", "wave")

    async def test_get_resources_meta(self, proxy_server):
        resources = await proxy_server.list_resources()
        wave_resource = next(r for r in resources if str(r.uri) == "resource://wave")
        assert wave_resource.title == "Wave"
        assert wave_resource.meta == {"fastmcp": {"tags": ["wave"]}}
        assert wave_resource.icons == [Icon(src="https://example.com/wave-icon.png")]

    async def test_list_resources_same_as_original(self, fastmcp_server, proxy_server):
        async with Client(fastmcp_server) as original_client:
            original = await original_client.list_resources()
        async with Client(proxy_server) as proxy_client:
            proxied = await proxy_client.list_resources()
        assert proxied == original

    async def test_read_resource(self, proxy_server: FastMCPProxy):
        async with Client(proxy_server) as client:
            result = await client.read_resource("resource://wave")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "👋"

    async def test_read_resource_same_as_original(self, fastmcp_server, proxy_server):
        async with Client(fastmcp_server) as client:
            result = await client.read_resource("resource://wave")
        async with Client(proxy_server) as client:
            proxy_result = await client.read_resource("resource://wave")
        assert proxy_result == result

    async def test_read_json_resource(self, proxy_server: FastMCPProxy):
        async with Client(proxy_server) as client:
            result = await client.read_resource("data://users")
        assert len(result) == 1
        assert isinstance(result[0], TextResourceContents)
        # The resource returns all users serialized as JSON
        users = json.loads(result[0].text)
        assert users == USERS

    async def test_proxy_returns_all_resource_contents(
        self, fastmcp_server, proxy_server
    ):
        """Test that proxy correctly returns all resource contents, not just the first one."""
        # Read from original server
        async with Client(fastmcp_server) as client:
            original_result = await client.read_resource("data://multi")

        # Read from proxy server
        async with Client(proxy_server) as client:
            proxy_result = await client.read_resource("data://multi")

        # Both should return the same number of contents
        assert len(original_result) == len(proxy_result)
        assert len(original_result) == 3

        # Verify all contents match
        for i, (original, proxied) in enumerate(zip(original_result, proxy_result)):
            assert isinstance(original, TextResourceContents)
            assert isinstance(proxied, TextResourceContents)
            assert original.text == proxied.text, f"Content {i} text mismatch"
            assert original.mime_type == proxied.mime_type, (
                f"Content {i} mimeType mismatch"
            )
            assert original.meta == proxied.meta, f"Content {i} meta mismatch"

        # Verify the contents are what we expect
        assert original_result[0].text == "First item"
        assert original_result[0].mime_type == "text/plain"
        assert original_result[1].text == '{"key": "value"}'
        assert original_result[1].mime_type == "application/json"
        assert original_result[2].text == "# Markdown\nContent"
        assert original_result[2].mime_type == "text/markdown"

    async def test_read_resource_returns_none_if_not_found(self, proxy_server):
        with pytest.raises(
            MCPError, match="Resource not found: 'resource://nonexistent'"
        ):
            async with Client(proxy_server) as client:
                await client.read_resource("resource://nonexistent")

    async def test_proxy_can_overwrite_proxied_resource(self, proxy_server):
        """
        Test that a resource defined on the proxy can overwrite the proxied resource with the same URI.
        """

        @proxy_server.resource(uri="resource://wave")
        def overwritten_wave() -> str:
            return "Overwritten wave! 🌊"

        async with Client(proxy_server) as client:
            result = await client.read_resource("resource://wave")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "Overwritten wave! 🌊"

    async def test_proxy_can_list_overwritten_resource(self, proxy_server):
        """
        Test that a resource defined on the proxy is listed instead of the proxied resource
        """

        @proxy_server.resource(uri="resource://wave", name="overwritten_wave")
        def overwritten_wave() -> str:
            return "Overwritten wave! 🌊"

        async with Client(proxy_server) as client:
            resources = await client.list_resources()
            wave_resource = next(
                r for r in resources if str(r.uri) == "resource://wave"
            )
            assert wave_resource.name == "overwritten_wave"


class TestResourceTemplates:
    async def test_get_resource_templates(self, proxy_server):
        templates = await proxy_server.list_resource_templates()
        assert [t.name for t in templates] == Contains("get_user")

    async def test_get_resource_templates_meta(self, proxy_server):
        templates = await proxy_server.list_resource_templates()
        get_user_template = next(
            t for t in templates if t.uri_template == "data://user/{user_id}"
        )
        assert get_user_template.title == "User Template"
        assert get_user_template.meta == {"fastmcp": {"tags": ["users"]}}
        assert get_user_template.icons == [
            Icon(src="https://example.com/user-icon.png")
        ]

    async def test_list_resource_templates_same_as_original(
        self, fastmcp_server, proxy_server
    ):
        async with Client(fastmcp_server) as original_client:
            result = await original_client.list_resource_templates()
        async with Client(proxy_server) as proxy_client:
            proxy_result = await proxy_client.list_resource_templates()
        assert proxy_result == result

    @pytest.mark.parametrize("id", [1, 2, 3])
    async def test_read_resource_template(self, proxy_server: FastMCPProxy, id: int):
        async with Client(proxy_server) as client:
            result = await client.read_resource(f"data://user/{id}")
        assert isinstance(result[0], TextResourceContents)
        assert json.loads(result[0].text) == USERS[id - 1]

    async def test_read_resource_template_same_as_original(
        self, fastmcp_server, proxy_server
    ):
        async with Client(fastmcp_server) as client:
            result = await client.read_resource("data://user/1")
        async with Client(proxy_server) as client:
            proxy_result = await client.read_resource("data://user/1")
        assert proxy_result == result

    async def test_proxy_template_returns_all_resource_contents(
        self, fastmcp_server, proxy_server
    ):
        """Test that proxy template correctly returns all resource contents."""
        # Read from original server
        async with Client(fastmcp_server) as client:
            original_result = await client.read_resource("data://multi/test123")

        # Read from proxy server
        async with Client(proxy_server) as client:
            proxy_result = await client.read_resource("data://multi/test123")

        # Both should return the same number of contents
        assert len(original_result) == len(proxy_result)
        assert len(original_result) == 2

        # Verify all contents match
        for i, (original, proxied) in enumerate(zip(original_result, proxy_result)):
            assert isinstance(original, TextResourceContents)
            assert isinstance(proxied, TextResourceContents)
            assert original.text == proxied.text, f"Content {i} text mismatch"
            assert original.mime_type == proxied.mime_type, (
                f"Content {i} mimeType mismatch"
            )

        # Verify the contents are what we expect
        assert original_result[0].text == "Item test123 - First"
        assert original_result[0].mime_type == "text/plain"
        assert original_result[1].text == '{"id": "test123", "status": "active"}'
        assert original_result[1].mime_type == "application/json"

    async def test_proxy_can_overwrite_proxied_resource_template(self, proxy_server):
        """
        Test that a resource template defined on the proxy can overwrite the proxied template with the same URI template.
        """

        @proxy_server.resource(uri="data://user/{user_id}", name="overwritten_get_user")
        def overwritten_get_user(user_id: str) -> str:
            return json.dumps(
                {
                    "id": user_id,
                    "name": "Overwritten User",
                    "active": True,
                    "extra": "data",
                }
            )

        async with Client(proxy_server) as client:
            result = await client.read_resource("data://user/1")
        assert isinstance(result[0], TextResourceContents)
        user_data = json.loads(result[0].text)
        assert user_data["name"] == "Overwritten User"
        assert user_data["extra"] == "data"

    async def test_proxy_can_list_overwritten_resource_template(self, proxy_server):
        """
        Test that a resource template defined on the proxy is listed instead of the proxied template
        """

        @proxy_server.resource(uri="data://user/{user_id}", name="overwritten_get_user")
        def overwritten_get_user(user_id: str) -> dict[str, Any]:
            return {"id": user_id, "name": "Overwritten User", "active": True}

        async with Client(proxy_server) as client:
            templates = await client.list_resource_templates()
            user_template = next(
                t for t in templates if t.uri_template == "data://user/{user_id}"
            )
            assert user_template.name == "overwritten_get_user"


class TestResourceTemplateQueryParams:
    """Resource templates with RFC 6570 {?param} query params work through proxy."""

    async def test_query_param_forwarded(self):
        remote = FastMCP("Remote")

        @remote.resource("data://{id}{?format}")
        def get_data(id: str, format: str = "json") -> str:
            return f"id={id} format={format}"

        proxy = create_proxy(Client(remote))
        async with Client(proxy) as client:
            result = await client.read_resource("data://123?format=xml")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "id=123 format=xml"

    async def test_query_param_default_used_when_omitted(self):
        remote = FastMCP("Remote")

        @remote.resource("data://{id}{?format}")
        def get_data(id: str, format: str = "json") -> str:
            return f"id={id} format={format}"

        proxy = create_proxy(Client(remote))
        async with Client(proxy) as client:
            result = await client.read_resource("data://123")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "id=123 format=json"

    async def test_multiple_query_params_forwarded(self):
        remote = FastMCP("Remote")

        @remote.resource("data://{id}{?limit,offset}")
        def get_data(id: str, limit: int = 10, offset: int = 0) -> str:
            return f"id={id} limit={limit} offset={offset}"

        proxy = create_proxy(Client(remote))
        async with Client(proxy) as client:
            result = await client.read_resource("data://abc?limit=5&offset=20")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "id=abc limit=5 offset=20"

    async def test_encoded_path_param_preserved(self):
        remote = FastMCP("Remote")

        @remote.resource("data://{id}")
        def get_data(id: str) -> str:
            return f"id={id}"

        proxy = create_proxy(Client(remote))
        async with Client(proxy) as client:
            result = await client.read_resource("data://a%2Fb")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "id=a/b"

    async def test_hyphenated_query_param_forwarded(self):
        remote = FastMCP("Remote")

        @remote.resource("data://{id}{?api-version}")
        def get_data(id: str, api_version: str = "v1") -> str:
            return f"id={id} api_version={api_version}"

        proxy = create_proxy(Client(remote))
        async with Client(proxy) as client:
            result = await client.read_resource("data://123?api-version=v2")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "id=123 api_version=v2"

    def test_same_name_in_path_and_query_is_rejected(self):
        remote = FastMCP("Remote")
        with pytest.raises(ValueError, match="must be optional"):

            @remote.resource("data://{id}{?id}")
            def get_data(id: str) -> str:
                return id

    async def test_hyphenated_query_param_not_double_encoded(self):
        remote = FastMCP("Remote")

        @remote.resource("data://{id}{?api-version}")
        def get_data(id: str, api_version: str = "v1") -> str:
            return f"id={id} api_version={api_version}"

        proxy = create_proxy(Client(remote))
        async with Client(proxy) as client:
            result = await client.read_resource("data://123?api-version=a%2Fb")
        assert isinstance(result[0], TextResourceContents)
        assert result[0].text == "id=123 api_version=a/b"


class TestPrompts:
    async def test_get_prompts_server_method(self, proxy_server: FastMCPProxy):
        prompts = await proxy_server.list_prompts()
        assert [p.name for p in prompts] == Contains("welcome")

    async def test_get_prompts_meta(self, proxy_server):
        prompts = await proxy_server.list_prompts()
        welcome_prompt = next(p for p in prompts if p.name == "welcome")
        assert welcome_prompt.title == "Welcome"
        assert welcome_prompt.meta == {"fastmcp": {"tags": ["welcome"]}}
        assert welcome_prompt.icons == [
            Icon(src="https://example.com/welcome-icon.png")
        ]

    async def test_list_prompts_same_as_original(self, fastmcp_server, proxy_server):
        async with Client(fastmcp_server) as client:
            result = await client.list_prompts()
        async with Client(proxy_server) as client:
            proxy_result = await client.list_prompts()
        assert proxy_result == result

    async def test_render_prompt_same_as_original(
        self, fastmcp_server: FastMCP, proxy_server: FastMCPProxy
    ):
        async with Client(fastmcp_server) as client:
            result = await client.get_prompt("welcome", {"name": "Alice"})
        async with Client(proxy_server) as client:
            proxy_result = await client.get_prompt("welcome", {"name": "Alice"})
        assert proxy_result == result

    async def test_render_prompt_calls_prompt(self, proxy_server):
        async with Client(proxy_server) as client:
            result = await client.get_prompt("welcome", {"name": "Alice"})
        assert result.messages[0].role == "user"
        assert isinstance(result.messages[0].content, TextContent)
        assert result.messages[0].content.text == "Welcome to FastMCP, Alice!"

    async def test_proxy_can_overwrite_proxied_prompt(self, proxy_server):
        """
        Test that a prompt defined on the proxy can overwrite the proxied prompt with the same name.
        """

        @proxy_server.prompt
        def welcome(name: str, extra: str = "friend") -> str:
            return f"Overwritten welcome, {name}! You are my {extra}."

        async with Client(proxy_server) as client:
            result = await client.get_prompt(
                "welcome", {"name": "Alice", "extra": "colleague"}
            )
        assert result.messages[0].role == "user"
        assert isinstance(result.messages[0].content, TextContent)
        assert (
            result.messages[0].content.text
            == "Overwritten welcome, Alice! You are my colleague."
        )

    async def test_proxy_can_list_overwritten_prompt(self, proxy_server):
        """
        Test that a prompt defined on the proxy is listed instead of the proxied prompt
        """

        @proxy_server.prompt
        def welcome(name: str, extra: str = "friend") -> str:
            return f"Overwritten welcome, {name}! You are my {extra}."

        async with Client(proxy_server) as client:
            prompts = await client.list_prompts()
            welcome_prompt = next(p for p in prompts if p.name == "welcome")
            # Check that the overwritten prompt has the additional 'extra' parameter
            param_names = [arg.name for arg in welcome_prompt.arguments or []]
            assert "extra" in param_names

    async def test_proxy_prompt_preserves_image_content(
        self, fastmcp_server: FastMCP, proxy_server: FastMCPProxy
    ):
        """Test that ProxyPrompt preserves ImageContent without lossy conversion."""
        async with Client(fastmcp_server) as client:
            result = await client.get_prompt("image_prompt")
        async with Client(proxy_server) as client:
            proxy_result = await client.get_prompt("image_prompt")

        # The proxy result should match the original exactly
        assert proxy_result == result
        # Verify the image content is preserved as ImageContent, not JSON text
        assert isinstance(proxy_result.messages[1].content, mcp_types.ImageContent)
        assert proxy_result.messages[1].content.data == "iVBORw0KGgoAAAANSUhEUg=="
        assert proxy_result.messages[1].content.mime_type == "image/png"


async def test_proxy_handles_multiple_concurrent_tasks_correctly(
    proxy_server: FastMCPProxy,
):
    results = {}

    async def get_and_store(name, coro):
        results[name] = await coro()

    async with create_task_group() as tg:
        tg.start_soon(get_and_store, "prompts", proxy_server.list_prompts)
        tg.start_soon(get_and_store, "resources", proxy_server.list_resources)
        tg.start_soon(get_and_store, "tools", proxy_server.list_tools)

    assert list(results) == Contains("resources", "prompts", "tools")
    assert [p.name for p in results["prompts"]] == Contains("welcome")
    assert [r.uri for r in results["resources"]] == Contains(
        AnyUrl("data://users"),
        AnyUrl("resource://wave"),
    )
    assert [r.name for r in results["resources"]] == Contains("get_users", "wave")
    assert [t.name for t in results["tools"]] == Contains(
        "greet", "add", "error_tool", "tool_without_description"
    )


class TestProxyComponentEnableDisable:
    """Test that enable/disable on proxy components guides users to server-level methods."""

    async def test_proxy_tool_enable_raises_not_implemented(self, proxy_server):
        """Test that enable() on proxy tools raises NotImplementedError."""
        tools = await proxy_server.list_tools()
        tool = next(t for t in tools if t.name == "greet")

        with pytest.raises(NotImplementedError, match="server.enable"):
            tool.enable()

    async def test_proxy_tool_disable_raises_not_implemented(self, proxy_server):
        """Test that disable() on proxy tools raises NotImplementedError."""
        tools = await proxy_server.list_tools()
        tool = next(t for t in tools if t.name == "greet")

        with pytest.raises(NotImplementedError, match="server.disable"):
            tool.disable()

    async def test_proxy_resource_enable_raises_not_implemented(self, proxy_server):
        """Test that enable() on proxy resources raises NotImplementedError."""
        resources = await proxy_server.list_resources()
        resource = next(r for r in resources if str(r.uri) == "resource://wave")

        with pytest.raises(NotImplementedError, match="server.enable"):
            resource.enable()

    async def test_proxy_resource_disable_raises_not_implemented(self, proxy_server):
        """Test that disable() on proxy resources raises NotImplementedError."""
        resources = await proxy_server.list_resources()
        resource = next(r for r in resources if str(r.uri) == "resource://wave")

        with pytest.raises(NotImplementedError, match="server.disable"):
            resource.disable()

    async def test_proxy_prompt_enable_raises_not_implemented(self, proxy_server):
        """Test that enable() on proxy prompts raises NotImplementedError."""
        prompts = await proxy_server.list_prompts()
        prompt = next(p for p in prompts if p.name == "welcome")

        with pytest.raises(NotImplementedError, match="server.enable"):
            prompt.enable()

    async def test_proxy_prompt_disable_raises_not_implemented(self, proxy_server):
        """Test that disable() on proxy prompts raises NotImplementedError."""
        prompts = await proxy_server.list_prompts()
        prompt = next(p for p in prompts if p.name == "welcome")

        with pytest.raises(NotImplementedError, match="server.disable"):
            prompt.disable()


class TestProxyProviderCache:
    """Tests for the ProxyProvider component list caching."""

    async def test_get_tool_uses_cached_list(self, fastmcp_server):
        """Calling call_tool should resolve from cache after an initial list."""
        provider = ProxyProvider(
            lambda: ProxyClient(FastMCPTransport(fastmcp_server)),
        )
        # Warm the cache via list
        tools = await provider.list_tools()
        assert any(t.name == "greet" for t in tools)

        # _get_tool should resolve from cache without calling _list_tools again
        with patch.object(
            provider, "_get_client", new_callable=AsyncMock
        ) as mock_client:
            tool = await provider._get_tool("greet")
            assert tool is not None
            assert tool.name == "greet"
            mock_client.assert_not_called()

    async def test_get_tool_fetches_on_cold_cache(self, fastmcp_server):
        """First _get_tool with no prior list should populate the cache."""
        provider = ProxyProvider(
            lambda: ProxyClient(FastMCPTransport(fastmcp_server)),
        )
        assert provider._tools_cache is None
        tool = await provider._get_tool("greet")
        assert tool is not None
        assert provider._tools_cache is not None

    async def test_cache_expires_after_ttl(self, fastmcp_server):
        """After TTL expires, _get_tool should re-fetch from the backend."""
        provider = ProxyProvider(
            lambda: ProxyClient(FastMCPTransport(fastmcp_server)),
            cache_ttl=0.0,
        )
        # Warm the cache
        await provider._list_tools()
        # With ttl=0 the cache is immediately stale, so _get_tool must re-fetch
        assert provider._tools_cache is not None
        original_ts = provider._tools_cache.timestamp

        time.sleep(0.05)

        await provider._get_tool("greet")
        assert provider._tools_cache.timestamp > original_ts

    async def test_list_tools_refreshes_cache(self, fastmcp_server):
        """Explicit list_tools always refreshes the cache timestamp."""
        provider = ProxyProvider(
            lambda: ProxyClient(FastMCPTransport(fastmcp_server)),
        )
        await provider._list_tools()
        first_ts = provider._tools_cache.timestamp  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]

        # Tiny sleep so monotonic clock advances
        time.sleep(0.05)

        await provider._list_tools()
        assert provider._tools_cache.timestamp > first_ts  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]

    async def test_cache_ttl_zero_disables_caching(self, fastmcp_server):
        """With cache_ttl=0, every _get_tool call should re-fetch."""
        provider = ProxyProvider(
            lambda: ProxyClient(FastMCPTransport(fastmcp_server)),
            cache_ttl=0.0,
        )
        # Each _get_tool call should trigger a fresh _list_tools
        call_count = 0
        original_list = provider._list_tools

        async def counting_list():
            nonlocal call_count
            call_count += 1
            return await original_list()

        with patch.object(provider, "_list_tools", side_effect=counting_list):
            await provider._get_tool("greet")
            await provider._get_tool("add")
        assert call_count == 2

    async def test_get_resource_uses_cache(self, fastmcp_server):
        """Resource lookups should also use the cache."""
        provider = ProxyProvider(
            lambda: ProxyClient(FastMCPTransport(fastmcp_server)),
        )
        await provider._list_resources()
        with patch.object(
            provider, "_get_client", new_callable=AsyncMock
        ) as mock_client:
            # Even if no resources match, the cache is used (no backend call)
            await provider._get_resource("config://app")
            mock_client.assert_not_called()

    async def test_call_tool_through_server_uses_cache(self, fastmcp_server):
        """End-to-end: calling a tool on a proxy server should only connect
        for the actual tool execution, not for tool resolution."""
        proxy = create_proxy(fastmcp_server)
        # Warm the cache by listing
        await proxy.list_tools()

        # Now call a tool — the provider's _list_tools should NOT be called
        # because the cache is warm. The connection happens only in ProxyTool.run.
        proxy_provider = next(
            p for p in proxy.providers if isinstance(p, ProxyProvider)
        )
        with patch.object(
            proxy_provider, "_list_tools", wraps=proxy_provider._list_tools
        ) as mock_list:
            result = await proxy.call_tool("greet", {"name": "Alice"})
            mock_list.assert_not_called()
        assert result.content[0].text == "Hello, Alice!"  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]


class TestProxySpanAttributes:
    """Regression tests for span attributes on un-renamed proxy components.

    A proxy component's private ``_backend_*`` field is only populated when
    the component is renamed via ``model_copy``. For un-renamed components it
    stays ``None``, so ``get_span_attributes()`` would emit ``None`` for the
    ``fastmcp.proxy.backend_*`` keys — which OpenTelemetry rejects with
    ``Invalid type NoneType for attribute ... value`` and drops, producing
    log spam on every proxied call.
    """

    async def test_proxy_tool_span_attributes_fall_back_to_name(self, proxy_server):
        proxy_provider = next(
            p for p in proxy_server.providers if isinstance(p, ProxyProvider)
        )
        tools = await proxy_provider._list_tools()
        assert tools, "expected the fixture to expose at least one tool"
        for tool in tools:
            attrs = tool.get_span_attributes()
            assert attrs["fastmcp.proxy.backend_name"] == tool.name
            assert all(v is not None for v in attrs.values()), (
                f"OpenTelemetry rejects None attribute values; got {attrs!r}"
            )

    async def test_proxy_resource_span_attributes_fall_back_to_uri(self, proxy_server):
        proxy_provider = next(
            p for p in proxy_server.providers if isinstance(p, ProxyProvider)
        )
        resources = await proxy_provider._list_resources()
        assert resources, "expected the fixture to expose at least one resource"
        for resource in resources:
            attrs = resource.get_span_attributes()
            assert attrs["fastmcp.proxy.backend_uri"] == str(resource.uri)
            assert all(v is not None for v in attrs.values()), (
                f"OpenTelemetry rejects None attribute values; got {attrs!r}"
            )

    async def test_proxy_template_span_attributes_fall_back_to_uri_template(
        self, proxy_server
    ):
        proxy_provider = next(
            p for p in proxy_server.providers if isinstance(p, ProxyProvider)
        )
        templates = await proxy_provider._list_resource_templates()
        assert templates, "expected the fixture to expose at least one template"
        for template in templates:
            attrs = template.get_span_attributes()
            assert attrs["fastmcp.proxy.backend_uri_template"] == template.uri_template
            assert all(v is not None for v in attrs.values()), (
                f"OpenTelemetry rejects None attribute values; got {attrs!r}"
            )

    async def test_proxy_prompt_span_attributes_fall_back_to_name(self, proxy_server):
        proxy_provider = next(
            p for p in proxy_server.providers if isinstance(p, ProxyProvider)
        )
        prompts = await proxy_provider._list_prompts()
        assert prompts, "expected the fixture to expose at least one prompt"
        for prompt in prompts:
            attrs = prompt.get_span_attributes()
            assert attrs["fastmcp.proxy.backend_name"] == prompt.name
            assert all(v is not None for v in attrs.values()), (
                f"OpenTelemetry rejects None attribute values; got {attrs!r}"
            )