prefecthq--fastmcp/tests/client/telemetry/test_client_tracing.py

"""Tests for client OpenTelemetry tracing."""

from __future__ import annotations

from collections.abc import AsyncGenerator

import pytest
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import SpanKind, StatusCode

from fastmcp import Client, FastMCP
from fastmcp.exceptions import ToolError


class TestClientToolTracing:
    """Tests for client tool call tracing."""

    async def test_call_tool_creates_span(self, trace_exporter: InMemorySpanExporter):
        server = FastMCP("test-server")

        @server.tool()
        def greet(name: str) -> str:
            return f"Hello, {name}!"

        client = Client(server)
        async with client:
            result = await client.call_tool("greet", {"name": "World"})
            assert "Hello, World!" in str(result)

        spans = trace_exporter.get_finished_spans()
        span_names = [s.name for s in spans]

        # Client should create "tools/call greet" span
        assert "tools/call greet" in span_names

    async def test_call_tool_span_attributes(
        self, trace_exporter: InMemorySpanExporter
    ):
        server = FastMCP("test-server")

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

        client = Client(server)
        async with client:
            await client.call_tool("add", {"a": 1, "b": 2})

        spans = trace_exporter.get_finished_spans()

        # Find client-side span (doesn't have fastmcp.server.name)
        client_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call add"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        assert client_span is not None
        assert client_span.attributes is not None
        # Standard MCP semantic conventions
        assert client_span.attributes["mcp.method.name"] == "tools/call"
        # gen_ai semantic conventions
        assert client_span.attributes["gen_ai.tool.name"] == "add"
        # RPC attributes must NOT be present
        assert "rpc.system" not in client_span.attributes
        assert "rpc.method" not in client_span.attributes
        # FastMCP-specific attributes
        assert client_span.attributes["fastmcp.component.key"] == "add"

    async def test_call_tool_error_caught_by_client_span(
        self, trace_exporter: InMemorySpanExporter
    ):
        """Tool error should be reflected on the client span via isError check."""
        server = FastMCP("test-server")

        @server.tool()
        def failing_tool() -> str:
            raise ValueError("boom")

        client = Client(server)
        async with client:
            with pytest.raises(ToolError):
                await client.call_tool("failing_tool", {})

        spans = trace_exporter.get_finished_spans()

        # Find the client span (from call_tool_mcp)
        client_spans = [
            s
            for s in spans
            if s.name == "tools/call failing_tool"
            and s.attributes is not None
            and "fastmcp.server.name" not in s.attributes
            and "fastmcp.component.key" in s.attributes
        ]

        # Exactly one client span should exist (no duplicate from call_tool)
        assert len(client_spans) == 1, (
            "There should be exactly one client span for call_tool"
        )

        error_span = client_spans[0]
        assert error_span.status.status_code == StatusCode.ERROR
        assert error_span.attributes is not None
        assert error_span.attributes["error.type"] == "tool_error"


class TestClientResourceTracing:
    """Tests for client resource read tracing."""

    async def test_read_resource_creates_span(
        self, trace_exporter: InMemorySpanExporter
    ):
        server = FastMCP("test-server")

        @server.resource("data://config")
        def get_config() -> str:
            return "config data"

        client = Client(server)
        async with client:
            result = await client.read_resource("data://config")
            assert "config data" in str(result)

        spans = trace_exporter.get_finished_spans()
        span_names = [s.name for s in spans]

        # Client should create "resources/read" span (URI in attributes, not name)
        assert "resources/read" in span_names

    async def test_read_resource_span_attributes(
        self, trace_exporter: InMemorySpanExporter
    ):
        server = FastMCP("test-server")

        @server.resource("data://config")
        def get_config() -> str:
            return "config value"

        client = Client(server)
        async with client:
            await client.read_resource("data://config")

        spans = trace_exporter.get_finished_spans()

        # Find client-side resource span (doesn't have fastmcp.server.name)
        client_span = next(
            (
                s
                for s in spans
                if s.name == "resources/read"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        assert client_span is not None
        assert client_span.attributes is not None
        # Standard MCP semantic conventions
        assert client_span.attributes["mcp.method.name"] == "resources/read"
        assert "data://" in str(client_span.attributes["mcp.resource.uri"])
        # RPC attributes must NOT be present
        assert "rpc.system" not in client_span.attributes
        assert "rpc.method" not in client_span.attributes
        # FastMCP-specific attributes
        # The URI may be normalized with trailing slash
        assert "data://" in str(client_span.attributes["fastmcp.component.key"])


class TestClientPromptTracing:
    """Tests for client prompt get tracing."""

    async def test_get_prompt_creates_span(self, trace_exporter: InMemorySpanExporter):
        server = FastMCP("test-server")

        @server.prompt()
        def greeting() -> str:
            return "Hello from prompt!"

        client = Client(server)
        async with client:
            result = await client.get_prompt("greeting")
            assert "Hello from prompt!" in str(result)

        spans = trace_exporter.get_finished_spans()
        span_names = [s.name for s in spans]

        # Client should create "prompts/get greeting" span
        assert "prompts/get greeting" in span_names

    async def test_get_prompt_span_attributes(
        self, trace_exporter: InMemorySpanExporter
    ):
        server = FastMCP("test-server")

        @server.prompt()
        def welcome(name: str) -> str:
            return f"Welcome, {name}!"

        client = Client(server)
        async with client:
            await client.get_prompt("welcome", {"name": "Test"})

        spans = trace_exporter.get_finished_spans()

        # Find client-side prompt span (doesn't have fastmcp.server.name)
        client_span = next(
            (
                s
                for s in spans
                if s.name == "prompts/get welcome"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        assert client_span is not None
        assert client_span.attributes is not None
        # Standard MCP semantic conventions
        assert client_span.attributes["mcp.method.name"] == "prompts/get"
        # gen_ai semantic conventions
        assert client_span.attributes["gen_ai.prompt.name"] == "welcome"
        # RPC attributes must NOT be present
        assert "rpc.system" not in client_span.attributes
        assert "rpc.method" not in client_span.attributes
        # FastMCP-specific attributes
        assert client_span.attributes["fastmcp.component.key"] == "welcome"


class TestClientServerSpanHierarchy:
    """Tests for span relationships between client and server."""

    async def test_client_and_server_spans_created(
        self, trace_exporter: InMemorySpanExporter
    ):
        """Both client and server should create spans for the same operation."""
        server = FastMCP("test-server")

        @server.tool()
        def echo(message: str) -> str:
            return message

        client = Client(server)
        async with client:
            await client.call_tool("echo", {"message": "test"})

        spans = trace_exporter.get_finished_spans()

        # Find client span (no fastmcp.server.name) and server span (has fastmcp.server.name)
        client_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call echo"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        server_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call echo"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        # Both spans should exist
        assert client_span is not None, "Client should create a span"
        assert client_span.attributes is not None
        assert server_span is not None, "Server should create a span"
        assert server_span.attributes is not None

        # Verify span kinds are correct
        assert client_span.kind == SpanKind.CLIENT, "Client span should be CLIENT kind"
        assert server_span.kind == SpanKind.SERVER, "Server span should be SERVER kind"

        # Verify the spans have different characteristics
        assert client_span.attributes["mcp.method.name"] == "tools/call"
        assert server_span.attributes["fastmcp.server.name"] == "test-server"

    async def test_trace_context_propagation(
        self, trace_exporter: InMemorySpanExporter
    ):
        """Server span should be a child of client span via trace context propagation."""
        server = FastMCP("test-server")

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

        client = Client(server)
        async with client:
            await client.call_tool("add", {"a": 1, "b": 2})

        spans = trace_exporter.get_finished_spans()

        # Find client span and server span
        client_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call add"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        server_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call add"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        assert client_span is not None, "Client span should exist"
        assert server_span is not None, "Server span should exist"

        # Verify trace context propagation: server span should be a descendant
        # of the client span, sharing the same trace_id. SDK v2 emits its own
        # intermediate spans (e.g. "MCP send tools/call add") between the FastMCP
        # client span and the server span, so the server span's parent is one of
        # those SDK spans rather than the client span directly. Assert descendant
        # relationship by walking the parent chain instead of a direct parent.
        assert server_span.context.trace_id == client_span.context.trace_id, (
            "Server and client spans should share the same trace_id"
        )

        spans_by_id = {s.context.span_id: s for s in spans}
        assert server_span.parent is not None, "Server span should have a parent"
        current = server_span
        found_client_ancestor = False
        while current.parent is not None:
            parent = spans_by_id.get(current.parent.span_id)
            if parent is None:
                break
            if parent.context.span_id == client_span.context.span_id:
                found_client_ancestor = True
                break
            current = parent
        assert found_client_ancestor, (
            "Server span should be a descendant of the client span"
        )


class TestClientErrorTracing:
    """Tests for client span creation during errors.

    Note: MCP protocol errors are returned as successful responses with error content,
    so client spans may not have ERROR status even when the operation fails. This is
    different from server-side where exceptions happen inside the span.

    The server-side span WILL have ERROR status because the exception occurs within
    the server's span context. The client span represents the successful MCP protocol
    round-trip, while application-level errors are communicated via the response.
    """

    async def test_call_tool_error_creates_spans(
        self, trace_exporter: InMemorySpanExporter
    ):
        """Both client and server spans should be created when tool fails."""
        server = FastMCP("test-server")

        @server.tool()
        def failing_tool() -> str:
            raise ValueError("Something went wrong")

        client = Client(server)
        async with client:
            with pytest.raises(ToolError):
                await client.call_tool("failing_tool", {})

        spans = trace_exporter.get_finished_spans()

        # Find client-side span
        client_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call failing_tool"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        # Find server-side span
        server_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call failing_tool"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        # Both spans should exist
        assert client_span is not None, "Client should create a span"
        assert server_span is not None, "Server should create a span"

        # Server span should have ERROR status (exception inside span)
        assert server_span.status.status_code == StatusCode.ERROR

    async def test_read_resource_error_creates_spans(
        self, trace_exporter: InMemorySpanExporter
    ):
        """Both client and server spans should be created when resource read fails."""
        server = FastMCP("test-server")

        @server.resource("data://fail")
        def failing_resource() -> str:
            raise ValueError("Resource error")

        client = Client(server)
        async with client:
            with pytest.raises(Exception):
                await client.read_resource("data://fail")

        spans = trace_exporter.get_finished_spans()

        # Find client-side span
        client_span = next(
            (
                s
                for s in spans
                if s.name == "resources/read"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        # Find server-side span
        server_span = next(
            (
                s
                for s in spans
                if s.name == "resources/read"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        # Both spans should exist
        assert client_span is not None, "Client should create a span"
        assert server_span is not None, "Server should create a span"

        # Server span should have ERROR status
        assert server_span.status.status_code == StatusCode.ERROR

    async def test_get_prompt_error_creates_spans(
        self, trace_exporter: InMemorySpanExporter
    ):
        """Both client and server spans should be created when prompt get fails."""
        server = FastMCP("test-server")

        @server.prompt()
        def failing_prompt() -> str:
            raise ValueError("Prompt error")

        client = Client(server)
        async with client:
            with pytest.raises(Exception):
                await client.get_prompt("failing_prompt", {})

        spans = trace_exporter.get_finished_spans()

        # Find client-side span
        client_span = next(
            (
                s
                for s in spans
                if s.name == "prompts/get failing_prompt"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        # Find server-side span
        server_span = next(
            (
                s
                for s in spans
                if s.name == "prompts/get failing_prompt"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        # Both spans should exist
        assert client_span is not None, "Client should create a span"
        assert server_span is not None, "Server should create a span"

        # Server span should have ERROR status
        assert server_span.status.status_code == StatusCode.ERROR

    async def test_call_nonexistent_tool_creates_spans(
        self, trace_exporter: InMemorySpanExporter
    ):
        """Both client and server spans should be created for nonexistent tool."""
        server = FastMCP("test-server")

        client = Client(server)
        async with client:
            with pytest.raises(Exception):
                await client.call_tool("nonexistent", {})

        spans = trace_exporter.get_finished_spans()

        # Find client-side span
        client_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call nonexistent"
                and s.attributes is not None
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        # Find server-side span
        server_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call nonexistent"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        # Both spans should exist
        assert client_span is not None, "Client should create a span"
        assert server_span is not None, "Server should create a span"

        # Server span should have ERROR status
        assert server_span.status.status_code == StatusCode.ERROR


class TestSessionIdOnSpans:
    """Tests for session ID capture on client and server spans.

    Session IDs are only available with HTTP transport (StreamableHttp).
    """

    @pytest.fixture
    async def http_server_url(self) -> AsyncGenerator[str, None]:
        """Start an HTTP server and return its URL."""
        from fastmcp.utilities.tests import run_server_async

        server = FastMCP("session-test-server")

        @server.tool()
        def echo(message: str) -> str:
            return message

        async with run_server_async(server) as url:
            yield url

    async def test_client_span_includes_session_id(
        self,
        trace_exporter: InMemorySpanExporter,
        http_server_url: str,
    ):
        """Client span should include session ID when using HTTP transport."""
        from fastmcp.client.transports import StreamableHttpTransport

        transport = StreamableHttpTransport(http_server_url)
        client = Client(transport=transport)
        async with client:
            await client.call_tool("echo", {"message": "test"})

        spans = trace_exporter.get_finished_spans()

        # Find the FastMCP client-side span. SDK v2 emits its own native OTel
        # spans (also lacking `fastmcp.server.name`), so select on the
        # `fastmcp.component.key` attribute that only our client_span sets.
        client_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call echo"
                and s.attributes is not None
                and "fastmcp.component.key" in s.attributes
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )

        assert client_span is not None, "Client should create a span"
        assert client_span.attributes is not None
        assert "mcp.session.id" in client_span.attributes
        assert client_span.attributes["mcp.session.id"] is not None

    async def test_server_span_includes_session_id(
        self,
        trace_exporter: InMemorySpanExporter,
        http_server_url: str,
    ):
        """Server span should include session ID when called via HTTP."""
        from fastmcp.client.transports import StreamableHttpTransport

        transport = StreamableHttpTransport(http_server_url)
        client = Client(transport=transport)
        async with client:
            await client.call_tool("echo", {"message": "test"})

        spans = trace_exporter.get_finished_spans()

        # Find server-side span
        server_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call echo"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        assert server_span is not None, "Server should create a span"
        assert server_span.attributes is not None
        assert "mcp.session.id" in server_span.attributes
        assert server_span.attributes["mcp.session.id"] is not None

    async def test_client_and_server_share_same_session_id(
        self,
        trace_exporter: InMemorySpanExporter,
        http_server_url: str,
    ):
        """Client and server spans should have the same session ID."""
        from fastmcp.client.transports import StreamableHttpTransport

        transport = StreamableHttpTransport(http_server_url)
        client = Client(transport=transport)
        async with client:
            await client.call_tool("echo", {"message": "test"})

        spans = trace_exporter.get_finished_spans()

        # Find both spans. Select the FastMCP client span on
        # `fastmcp.component.key` to avoid matching SDK v2's native OTel span.
        client_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call echo"
                and s.attributes is not None
                and "fastmcp.component.key" in s.attributes
                and "fastmcp.server.name" not in s.attributes
                and "fastmcp.component.key" in s.attributes
            ),
            None,
        )
        server_span = next(
            (
                s
                for s in spans
                if s.name == "tools/call echo"
                and s.attributes is not None
                and "fastmcp.server.name" in s.attributes
            ),
            None,
        )

        assert client_span is not None
        assert client_span.attributes is not None
        assert server_span is not None
        assert server_span.attributes is not None

        # Both should have session IDs and they should match
        client_session = client_span.attributes.get("mcp.session.id")
        server_session = server_span.attributes.get("mcp.session.id")

        assert client_session is not None, "Client span should have session ID"
        assert server_session is not None, "Server span should have session ID"
        assert client_session == server_session, (
            "Client and server should share the same session ID"
        )