diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py index c2004be7b..b0ec7cde3 100644 --- a/src/fastmcp/client/client.py +++ b/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() diff --git a/src/fastmcp/client/mixins/prompts.py b/src/fastmcp/client/mixins/prompts.py index de503b448..168f12037 100644 --- a/src/fastmcp/client/mixins/prompts.py +++ b/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 diff --git a/src/fastmcp/client/mixins/resources.py b/src/fastmcp/client/mixins/resources.py index 3cc845d50..542e20cd6 100644 --- a/src/fastmcp/client/mixins/resources.py +++ b/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 diff --git a/src/fastmcp/client/mixins/tools.py b/src/fastmcp/client/mixins/tools.py index c66c6e1f7..1a986561c 100644 --- a/src/fastmcp/client/mixins/tools.py +++ b/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 diff --git a/src/fastmcp/telemetry.py b/src/fastmcp/telemetry.py index 0965b8b71..b881a0632 100644 --- a/src/fastmcp/telemetry.py +++ b/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", diff --git a/tests/client/telemetry/test_client_propagation.py b/tests/client/telemetry/test_client_propagation.py new file mode 100644 index 000000000..4acde206d --- /dev/null +++ b/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"]) diff --git a/tests/telemetry/test_module.py b/tests/telemetry/test_module.py index c760af5cf..23359bf43 100644 --- a/tests/telemetry/test_module.py +++ b/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):