Propagate trace context on client initialize and list methods

🤖 Generated with Codex

src/fastmcp/client/client.py

@@ -17,6 +17,12 @@ import httpx
 import mcp.types
 from exceptiongroup import catch
 from mcp import ClientSession, McpError
+from mcp.client.session import (
+    SUPPORTED_PROTOCOL_VERSIONS,
+    _default_elicitation_callback,
+    _default_list_roots_callback,
+    _default_sampling_callback,
+)
 from mcp.types import GetTaskResult, TaskStatusNotification
 from pydantic import AnyUrl
 
@@ -51,8 +57,10 @@ from fastmcp.client.tasks import (
     TaskNotificationHandler,
     ToolTask,
 )
+from fastmcp.client.telemetry import client_span
 from fastmcp.mcp_config import MCPConfig
 from fastmcp.server import FastMCP
+from fastmcp.telemetry import inject_trace_context
 from fastmcp.utilities.exceptions import get_catch_handlers
 from fastmcp.utilities.logging import get_logger
 from fastmcp.utilities.timeout import (
@@ -518,12 +526,83 @@ class Client(
             timeout = normalize_timeout_to_seconds(timeout)
 
         try:
-            with anyio.fail_after(timeout):
-                self._session_state.initialize_result = await self.session.initialize()
-                return self._session_state.initialize_result
+            with client_span(
+                "initialize",
+                "initialize",
+                "",
+                session_id=self.transport.get_session_id(),
+            ):
+                with anyio.fail_after(timeout):
+                    propagated_meta = inject_trace_context()
+                    if propagated_meta is None:
+                        self._session_state.initialize_result = (
+                            await self.session.initialize()
+                        )
+                    else:
+                        self._session_state.initialize_result = (
+                            await self._initialize_session_with_meta(propagated_meta)
+                        )
+                    return self._session_state.initialize_result
         except TimeoutError as e:
             raise RuntimeError("Failed to initialize server session") from e
 
+    async def _initialize_session_with_meta(
+        self,
+        meta: dict[str, Any],
+    ) -> mcp.types.InitializeResult:
+        """Initialize the MCP session while propagating MCP ``_meta`` fields."""
+        session = self.session
+
+        sampling = (
+            (session._sampling_capabilities or mcp.types.SamplingCapability())
+            if session._sampling_callback is not _default_sampling_callback
+            else None
+        )
+        elicitation = (
+            mcp.types.ElicitationCapability(
+                form=mcp.types.FormElicitationCapability(),
+                url=mcp.types.UrlElicitationCapability(),
+            )
+            if session._elicitation_callback is not _default_elicitation_callback
+            else None
+        )
+        roots = (
+            mcp.types.RootsCapability(listChanged=True)
+            if session._list_roots_callback is not _default_list_roots_callback
+            else None
+        )
+
+        result = await session.send_request(
+            mcp.types.ClientRequest(
+                mcp.types.InitializeRequest(
+                    params=mcp.types.InitializeRequestParams(
+                        protocolVersion=mcp.types.LATEST_PROTOCOL_VERSION,
+                        capabilities=mcp.types.ClientCapabilities(
+                            sampling=sampling,
+                            elicitation=elicitation,
+                            experimental=None,
+                            roots=roots,
+                            tasks=session._task_handlers.build_capability(),
+                        ),
+                        clientInfo=session._client_info,
+                        _meta=meta,  # type: ignore[unknown-argument]  # pydantic alias  # ty:ignore[unknown-argument]
+                    )
+                )
+            ),
+            mcp.types.InitializeResult,
+        )
+
+        if result.protocolVersion not in SUPPORTED_PROTOCOL_VERSIONS:
+            raise RuntimeError(
+                f"Unsupported protocol version from the server: {result.protocolVersion}"
+            )
+
+        session._server_capabilities = result.capabilities
+        await session.send_notification(
+            mcp.types.ClientNotification(mcp.types.InitializedNotification())
+        )
+        return result
+
     async def __aenter__(self):
         return await self._connect()
 

src/fastmcp/client/mixins/prompts.py

@@ -57,8 +57,16 @@ class ClientPromptsMixin:
         ):
             logger.debug(f"[{self.name}] called list_prompts")
 
+            propagated_meta = inject_trace_context()
+            params = None
+            if cursor is not None or propagated_meta is not None:
+                params = mcp.types.PaginatedRequestParams(
+                    cursor=cursor,
+                    _meta=propagated_meta,  # type: ignore[unknown-argument]  # pydantic alias  # ty:ignore[unknown-argument]
+                )
+
             result = await self._await_with_session_monitoring(
-                self.session.list_prompts(cursor=cursor)
+                self.session.list_prompts(params=params)
             )
             return result
 

src/fastmcp/client/mixins/resources.py

@@ -56,8 +56,16 @@ class ClientResourcesMixin:
         ):
             logger.debug(f"[{self.name}] called list_resources")
 
+            propagated_meta = inject_trace_context()
+            params = None
+            if cursor is not None or propagated_meta is not None:
+                params = mcp.types.PaginatedRequestParams(
+                    cursor=cursor,
+                    _meta=propagated_meta,  # type: ignore[unknown-argument]  # pydantic alias  # ty:ignore[unknown-argument]
+                )
+
             result = await self._await_with_session_monitoring(
-                self.session.list_resources(cursor=cursor)
+                self.session.list_resources(params=params)
             )
             return result
 
@@ -132,8 +140,16 @@ class ClientResourcesMixin:
         ):
             logger.debug(f"[{self.name}] called list_resource_templates")
 
+            propagated_meta = inject_trace_context()
+            params = None
+            if cursor is not None or propagated_meta is not None:
+                params = mcp.types.PaginatedRequestParams(
+                    cursor=cursor,
+                    _meta=propagated_meta,  # type: ignore[unknown-argument]  # pydantic alias  # ty:ignore[unknown-argument]
+                )
+
             result = await self._await_with_session_monitoring(
-                self.session.list_resource_templates(cursor=cursor)
+                self.session.list_resource_templates(params=params)
             )
             return result
 

src/fastmcp/client/mixins/tools.py

@@ -61,8 +61,16 @@ class ClientToolsMixin:
         ):
             logger.debug(f"[{self.name}] called list_tools")
 
+            propagated_meta = inject_trace_context()
+            params = None
+            if cursor is not None or propagated_meta is not None:
+                params = mcp.types.PaginatedRequestParams(
+                    cursor=cursor,
+                    _meta=propagated_meta,  # type: ignore[unknown-argument]  # pydantic alias  # ty:ignore[unknown-argument]
+                )
+
             result = await self._await_with_session_monitoring(
-                self.session.list_tools(cursor=cursor)
+                self.session.list_tools(params=params)
             )
             return result
 

src/fastmcp/telemetry.py

@@ -33,6 +33,7 @@ INSTRUMENTATION_NAME = "fastmcp"
 
 TRACE_PARENT_KEY = "traceparent"
 TRACE_STATE_KEY = "tracestate"
+BAGGAGE_KEY = "baggage"
 
 
 def get_tracer(version: str | None = None) -> Tracer:
@@ -67,6 +68,8 @@ def inject_trace_context(
         trace_meta[TRACE_PARENT_KEY] = carrier["traceparent"]
     if "tracestate" in carrier:
         trace_meta[TRACE_STATE_KEY] = carrier["tracestate"]
+    if BAGGAGE_KEY in carrier:
+        trace_meta[BAGGAGE_KEY] = carrier[BAGGAGE_KEY]
 
     if trace_meta:
         return {**(meta or {}), **trace_meta}
@@ -112,6 +115,7 @@ def extract_trace_context(meta: dict[str, Any] | None) -> Context:
 
 
 __all__ = [
+    "BAGGAGE_KEY",
     "INSTRUMENTATION_NAME",
     "TRACE_PARENT_KEY",
     "TRACE_STATE_KEY",

tests/client/telemetry/test_client_propagation.py

@@ -0,0 +1,180 @@
+"""Tests for client trace-context propagation on initialize and list methods."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+import mcp.types as mt
+import pytest
+from opentelemetry import baggage, trace
+from opentelemetry import context as otel_context
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from fastmcp import Client, FastMCP
+
+
+def build_test_server() -> FastMCP:
+    """Create a server with one component of each type."""
+    server = FastMCP("test-server")
+
+    @server.tool()
+    def greet(name: str = "world") -> str:
+        return f"Hello, {name}!"
+
+    @server.resource("data://config")
+    def get_config() -> str:
+        return "config"
+
+    @server.resource("users://{user_id}/profile")
+    def get_profile(user_id: str) -> str:
+        return f"profile {user_id}"
+
+    @server.prompt()
+    def greeting() -> str:
+        return "Hello!"
+
+    return server
+
+
+def get_meta_dict(request: Any) -> dict[str, Any] | None:
+    """Normalize ``request.params.meta`` to a plain dict for assertions."""
+    params = getattr(request, "params", None)
+    meta = getattr(params, "meta", None)
+    if meta is None:
+        return None
+    return meta.model_dump(exclude_none=True)
+
+
+class TestClientPropagation:
+    async def test_initialize_creates_span_and_propagates_meta(
+        self,
+        monkeypatch,
+        trace_exporter: InMemorySpanExporter,
+    ):
+        server = build_test_server()
+        client = Client(server, auto_initialize=False)
+        tracer = trace.get_tracer("external")
+
+        async with client:
+            captured_requests: list[Any] = []
+            original_send_request = client.session.send_request
+
+            async def wrapped_send_request(request: Any, result_type: Any) -> Any:
+                captured_requests.append(request.root)
+                return await original_send_request(request, result_type)
+
+            monkeypatch.setattr(client.session, "send_request", wrapped_send_request)
+
+            baggage_token = otel_context.attach(baggage.set_baggage("tenant", "acme"))
+            try:
+                with tracer.start_as_current_span("external-parent"):
+                    await client.initialize()
+            finally:
+                otel_context.detach(baggage_token)
+
+        spans = trace_exporter.get_finished_spans()
+        initialize_span = next(
+            (
+                span
+                for span in spans
+                if span.name == "initialize"
+                and span.attributes is not None
+                and "fastmcp.server.name" not in span.attributes
+            ),
+            None,
+        )
+        assert initialize_span is not None
+        assert initialize_span.attributes is not None
+        assert initialize_span.attributes["mcp.method.name"] == "initialize"
+
+        initialize_request = next(
+            request
+            for request in captured_requests
+            if isinstance(request, mt.InitializeRequest)
+        )
+        captured_meta = get_meta_dict(initialize_request)
+        assert captured_meta is not None
+        assert captured_meta["traceparent"].split("-")[2] == format(
+            initialize_span.get_span_context().span_id,
+            "016x",
+        )
+        assert "tenant=acme" in str(captured_meta["baggage"])
+
+    @pytest.mark.parametrize(
+        ("request_type", "operation", "expected_span_name"),
+        [
+            (mt.ListToolsRequest, lambda client: client.list_tools(), "tools/list"),
+            (
+                mt.ListResourcesRequest,
+                lambda client: client.list_resources(),
+                "resources/list",
+            ),
+            (
+                mt.ListResourceTemplatesRequest,
+                lambda client: client.list_resource_templates(),
+                "resources/templates/list",
+            ),
+            (
+                mt.ListPromptsRequest,
+                lambda client: client.list_prompts(),
+                "prompts/list",
+            ),
+        ],
+    )
+    async def test_list_methods_propagate_meta(
+        self,
+        request_type: type[Any],
+        operation: Callable[[Client], Awaitable[Any]],
+        expected_span_name: str,
+        monkeypatch,
+        trace_exporter: InMemorySpanExporter,
+    ):
+        server = build_test_server()
+        client = Client(server, auto_initialize=False)
+        tracer = trace.get_tracer("external")
+
+        async with client:
+            await client.initialize()
+            captured_requests: list[Any] = []
+            original_send_request = client.session.send_request
+
+            async def wrapped_send_request(request: Any, result_type: Any) -> Any:
+                captured_requests.append(request.root)
+                return await original_send_request(request, result_type)
+
+            monkeypatch.setattr(client.session, "send_request", wrapped_send_request)
+            trace_exporter.clear()
+
+            baggage_token = otel_context.attach(baggage.set_baggage("tenant", "acme"))
+            try:
+                with tracer.start_as_current_span("external-parent"):
+                    await operation(client)
+            finally:
+                otel_context.detach(baggage_token)
+
+        spans = trace_exporter.get_finished_spans()
+        client_span = next(
+            (
+                span
+                for span in spans
+                if span.name == expected_span_name
+                and span.attributes is not None
+                and "fastmcp.server.name" not in span.attributes
+            ),
+            None,
+        )
+        assert client_span is not None
+
+        request = next(
+            captured_request
+            for captured_request in captured_requests
+            if isinstance(captured_request, request_type)
+        )
+        captured_meta = get_meta_dict(request)
+        assert captured_meta is not None
+        assert captured_meta["traceparent"].split("-")[2] == format(
+            client_span.get_span_context().span_id,
+            "016x",
+        )
+        assert "tenant=acme" in str(captured_meta["baggage"])

tests/telemetry/test_module.py

@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
-from opentelemetry import trace
+from opentelemetry import baggage, trace
+from opentelemetry import context as otel_context
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
 from fastmcp.server.telemetry import get_auth_span_attributes
 from fastmcp.telemetry import (
+    BAGGAGE_KEY,
     INSTRUMENTATION_NAME,
     TRACE_PARENT_KEY,
     extract_trace_context,
@@ -50,6 +52,19 @@ class TestInjectTraceContext:
         assert TRACE_PARENT_KEY in meta
         assert meta[TRACE_PARENT_KEY].startswith("00-")
 
+    def test_injects_baggage(self, trace_exporter: InMemorySpanExporter):
+        tracer = get_tracer()
+        baggage_token = otel_context.attach(baggage.set_baggage("userId", "alice"))
+        try:
+            with tracer.start_as_current_span("test"):
+                meta = inject_trace_context()
+        finally:
+            otel_context.detach(baggage_token)
+
+        assert meta is not None
+        assert BAGGAGE_KEY in meta
+        assert "userId=alice" in str(meta[BAGGAGE_KEY])
+
 
 class TestExtractTraceContext:
     def test_bare_traceparent(self, trace_exporter: InMemorySpanExporter):