refactor(MCP): Pass MCP call context to the API as W3C Baggage

Write flagsmith.client.name/version as span attributes directly, and
pass all of the current span's attributes — MCP method, tool, session
and client identity — as W3C Baggage on outbound Flagsmith API
requests via an httpx event hook scoped to the API client only.
Replaces the baggage-attaching middleware: process context is never
touched, and the API receives the full call context whenever tracing
is enabled.

beep boop

Author: khvn26

mcp/src/flagsmith_mcp/server.py

@@ -16,7 +16,7 @@
 from flagsmith_mcp.events import EventLoggingMiddleware
 from flagsmith_mcp.metrics import PrometheusMiddleware
 from flagsmith_mcp.oauth import FlagsmithResourceAuth
-from flagsmith_mcp.telemetry import BaggageMiddleware, setup_telemetry
+from flagsmith_mcp.telemetry import propagate_span_attributes, setup_telemetry
 
 ROUTE_MAPS = [
     RouteMap(tags={"mcp"}, mcp_type=MCPType.TOOL),
@@ -61,9 +61,11 @@ def create_server(settings: config.Settings) -> FastMCP[None]:
     api_client = httpx.AsyncClient(
         base_url=settings.flagsmith_api_url,
         auth=FlagsmithAuth(settings.flagsmith_api_token),
+        event_hooks={"request": [propagate_span_attributes]},
     )
     # Instrument only the Flagsmith API client: emit a span per upstream
-    # call and propagate W3C trace context and baggage to the API.
+    # call and propagate W3C trace context; the event hook passes the MCP
+    # call context to the API as W3C Baggage.
     HTTPXClientInstrumentor().instrument_client(api_client)
     server = FastMCP.from_openapi(
         openapi_spec=_fetch_spec(),
@@ -77,7 +79,6 @@ def create_server(settings: config.Settings) -> FastMCP[None]:
 
     server.add_middleware(PrometheusMiddleware())
     server.add_middleware(EventLoggingMiddleware())
-    server.add_middleware(BaggageMiddleware())
 
     @server.custom_route("/health", methods=["GET"])
     async def health(request: Request) -> PlainTextResponse:

mcp/src/flagsmith_mcp/telemetry.py

@@ -1,19 +1,15 @@
-from typing import Any
-
-import mcp.types as mt
+import httpx
 from common.core.logging import setup_logging
 from common.core.otel import (
     add_otel_trace_context,
     build_otel_log_provider,
     build_tracer_provider,
     make_structlog_otel_processor,
 )
-from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
-from fastmcp.tools.base import ToolResult
 from opentelemetry import baggage, trace
-from opentelemetry import context as otel_context
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
 from opentelemetry.context import Context
-from opentelemetry.sdk.trace import Span, SpanProcessor
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from structlog.typing import Processor
 
 from flagsmith_mcp import config
@@ -22,59 +18,24 @@
 APPLICATION_LOGGERS = ["flagsmith_mcp", "fastmcp", "mcp"]
 
 
-def is_flagsmith_baggage_key(key: str) -> bool:
-    return key.startswith("flagsmith.")
-
-
-class FlagsmithBaggageSpanProcessor(SpanProcessor):
-    """Copy flagsmith.* baggage entries onto started spans.
-
-    The off-the-shelf opentelemetry-processor-baggage reads baggage from
-    the span's parent context, but FastMCP starts its server spans from a
-    context extracted from the request _meta, which carries no baggage.
-    Read the current context, where middleware attached the entries.
-    """
+class ClientInfoSpanProcessor(SpanProcessor):
+    """Annotate started spans with the MCP client identity."""
 
     def on_start(self, span: Span, parent_context: Context | None = None) -> None:
-        for key, value in baggage.get_all().items():
-            if is_flagsmith_baggage_key(key):
-                span.set_attribute(key, str(value))
+        if (client_info := get_client_info()) is not None:
+            span.set_attribute("flagsmith.client.name", client_info.name)
+            span.set_attribute("flagsmith.client.version", client_info.version)
 
 
-class BaggageMiddleware(Middleware):
-    """Attach the client identity and tool name as W3C Baggage: the single
-    source both for span attributes (via BaggageSpanProcessor) and for
-    propagation to the Flagsmith API by the instrumented upstream client."""
-
-    async def on_request(
-        self,
-        context: MiddlewareContext[mt.Request[Any, Any]],
-        call_next: CallNext[mt.Request[Any, Any], Any],
-    ) -> Any:
-        if (client_info := get_client_info()) is None:
-            return await call_next(context)
-        ctx = baggage.set_baggage("flagsmith.client.name", client_info.name)
-        ctx = baggage.set_baggage(
-            "flagsmith.client.version", client_info.version, context=ctx
-        )
-        token = otel_context.attach(ctx)
-        try:
-            return await call_next(context)
-        finally:
-            otel_context.detach(token)
-
-    async def on_call_tool(
-        self,
-        context: MiddlewareContext[mt.CallToolRequestParams],
-        call_next: CallNext[mt.CallToolRequestParams, ToolResult],
-    ) -> ToolResult:
-        token = otel_context.attach(
-            baggage.set_baggage("flagsmith.tool.name", context.message.name)
-        )
-        try:
-            return await call_next(context)
-        finally:
-            otel_context.detach(token)
+async def propagate_span_attributes(request: httpx.Request) -> None:
+    span = trace.get_current_span()
+    if not isinstance(span, ReadableSpan):
+        return
+    ctx: Context | None = None
+    for key, value in (span.attributes or {}).items():
+        ctx = baggage.set_baggage(key, str(value), context=ctx)
+    if ctx is not None:
+        W3CBaggagePropagator().inject(request.headers, context=ctx)
 
 
 def setup_telemetry(settings: config.Settings) -> None:
@@ -97,7 +58,7 @@ def setup_telemetry(settings: config.Settings) -> None:
             endpoint=f"{endpoint}/v1/traces",
             service_name=settings.otel_service_name,
         )
-        tracer_provider.add_span_processor(FlagsmithBaggageSpanProcessor())
+        tracer_provider.add_span_processor(ClientInfoSpanProcessor())
         trace.set_tracer_provider(tracer_provider)
     setup_logging(
         log_level=settings.log_level,

mcp/tests/integration/conftest.py

@@ -16,7 +16,7 @@
 
 from flagsmith_mcp import config, constants
 from flagsmith_mcp import server as server_module
-from flagsmith_mcp.telemetry import FlagsmithBaggageSpanProcessor
+from flagsmith_mcp.telemetry import ClientInfoSpanProcessor
 
 HTTPClientFactoryFixture = Callable[[FastMCP], AsyncIterator[httpx.AsyncClient]]
 
@@ -28,7 +28,7 @@ def span_exporter() -> InMemorySpanExporter:
     exporter = InMemorySpanExporter()
     provider = TracerProvider()
     provider.add_span_processor(SimpleSpanProcessor(exporter))
-    provider.add_span_processor(FlagsmithBaggageSpanProcessor())
+    provider.add_span_processor(ClientInfoSpanProcessor())
     trace.set_tracer_provider(provider)
     return exporter
 

mcp/tests/integration/test_spans.py

@@ -1,3 +1,5 @@
+from urllib.parse import unquote
+
 from fastmcp import Client
 from fastmcp.client.transports import FastMCPTransport
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
@@ -29,15 +31,15 @@ async def test_spans__tool_call__annotated_with_client_identity(
     assert span.attributes is not None
     assert {
         "gen_ai.tool.name": "list_environments",
-        "flagsmith.tool.name": "list_environments",
         "flagsmith.client.name": "mcp",
         "flagsmith.client.version": "0.1.0",
     }.items() <= dict(span.attributes).items()
 
 
-async def test_spans__tool_call__upstream_request_carries_baggage(
+async def test_spans__tool_call__upstream_request_carries_span_attribute_baggage(
     client: Client[FastMCPTransport],
     respx_mock: MockRouter,
+    finished_spans: InMemorySpanExporter,
 ) -> None:
     # Given
     route = respx_mock.get("https://api.flagsmith.com/environments/")
@@ -46,9 +48,17 @@ async def test_spans__tool_call__upstream_request_carries_baggage(
     # When
     await client.call_tool("list_environments", {})
 
-    # Then the instrumented API client propagated W3C Baggage
-    assert route.calls.last.request.headers["baggage"] == (
-        "flagsmith.client.name=mcp,"
-        "flagsmith.client.version=0.1.0,"
-        "flagsmith.tool.name=list_environments"
-    )
+    # Then the server span's attributes reached the API as W3C Baggage
+    entries = {
+        key: unquote(value)
+        for key, value in (
+            entry.split("=", 1)
+            for entry in route.calls.last.request.headers["baggage"].split(",")
+        )
+    }
+    assert {
+        "mcp.method.name": "tools/call",
+        "gen_ai.tool.name": "list_environments",
+        "flagsmith.client.name": "mcp",
+        "flagsmith.client.version": "0.1.0",
+    }.items() <= entries.items()

mcp/tests/unit/test_event_logging.py

@@ -10,22 +10,22 @@ def test_get_client_info__outside_request_context__returns_none() -> None:
     assert events.get_client_info() is None
 
 
-def test_get_client_info__uninitialised_session__returns_none(
+def test_get_client_info__no_client_params__returns_none(
     mocker: MockerFixture,
 ) -> None:
-    # Given a session that has not completed initialize
+    # Given a session that has not yet completed initialize
     context_mock = mocker.patch.object(events, "get_context", autospec=True)
     context_mock.return_value.session.client_params = None
 
     # When / Then
     assert events.get_client_info() is None
 
 
-async def test_event_logging_middleware__uninitialised_session__empty_client_identity(
+async def test_event_logging_middleware__no_client_info__empty_client_identity(
     mocker: MockerFixture,
     log: StructuredLogCapture,
 ) -> None:
-    # Given a tool call outside an initialised session
+    # Given
     middleware = events.EventLoggingMiddleware()
     context = mocker.Mock()
     context.message.name = "list_environments"

mcp/tests/unit/test_telemetry.py

@@ -1,8 +1,9 @@
 import os
 
+import httpx
 from common.core.otel import add_otel_trace_context
-from opentelemetry import baggage
-from opentelemetry import context as otel_context
+from opentelemetry import trace
+from opentelemetry.sdk.trace import ReadableSpan
 from pytest_mock import MockerFixture
 
 from flagsmith_mcp import config, telemetry
@@ -77,7 +78,7 @@ def test_setup_telemetry__otlp_endpoint__exports_logs_and_traces(
     [span_processor] = (
         build_tracer_provider_mock.return_value.add_span_processor.call_args.args
     )
-    assert isinstance(span_processor, telemetry.FlagsmithBaggageSpanProcessor)
+    assert isinstance(span_processor, telemetry.ClientInfoSpanProcessor)
     setup_logging_mock.assert_called_once_with(
         log_level="DEBUG",
         log_format="json",
@@ -89,63 +90,43 @@ def test_setup_telemetry__otlp_endpoint__exports_logs_and_traces(
     )
 
 
-async def test_baggage_middleware__uninitialised_session__tool_name_baggage_only(
+def test_client_info_span_processor__outside_request_context__no_attributes(
     mocker: MockerFixture,
 ) -> None:
-    # Given a tool call outside an initialised session
-    middleware = telemetry.BaggageMiddleware()
-    context = mocker.Mock()
-    context.message.name = "list_environments"
-    seen_baggage: dict[str, object] = {}
-
-    async def record_baggage(ctx: object) -> None:
-        seen_baggage.update(baggage.get_all())
-
-    call_next = mocker.AsyncMock(side_effect=record_baggage)
+    # Given no MCP request context
+    span = mocker.Mock()
 
     # When
-    await middleware.on_call_tool(context, call_next)
+    telemetry.ClientInfoSpanProcessor().on_start(span)
 
     # Then
-    assert seen_baggage == {"flagsmith.tool.name": "list_environments"}
-    assert baggage.get_all() == {}
+    span.set_attribute.assert_not_called()
 
 
-async def test_baggage_middleware__uninitialised_session__request_baggage_untouched(
-    mocker: MockerFixture,
-) -> None:
-    # Given a request outside an initialised session
-    middleware = telemetry.BaggageMiddleware()
-    seen_baggage: dict[str, object] = {}
-
-    async def record_baggage(ctx: object) -> None:
-        seen_baggage.update(baggage.get_all())
-
-    call_next = mocker.AsyncMock(side_effect=record_baggage)
+async def test_propagate_span_attributes__no_recording_span__headers_untouched() -> (
+    None
+):
+    # Given no SDK span in the current context
+    request = httpx.Request("GET", "https://api.flagsmith.com/")
 
     # When
-    await middleware.on_request(mocker.Mock(), call_next)
+    await telemetry.propagate_span_attributes(request)
 
     # Then
-    assert seen_baggage == {}
+    assert "baggage" not in request.headers
 
 
-def test_flagsmith_baggage_span_processor__foreign_baggage__not_copied(
+async def test_propagate_span_attributes__span_without_attributes__headers_untouched(
     mocker: MockerFixture,
 ) -> None:
-    # Given baggage with flagsmith and foreign entries
-    span = mocker.Mock()
-    ctx = baggage.set_baggage("other.key", "x")
-    ctx = baggage.set_baggage("flagsmith.tool.name", "list_environments", context=ctx)
-    token = otel_context.attach(ctx)
+    # Given a recording span with no attributes
+    span = mocker.Mock(spec=ReadableSpan)
+    span.attributes = {}
+    mocker.patch.object(trace, "get_current_span", return_value=span)
 
     # When
-    try:
-        telemetry.FlagsmithBaggageSpanProcessor().on_start(span)
-    finally:
-        otel_context.detach(token)
+    request = httpx.Request("GET", "https://api.flagsmith.com/")
+    await telemetry.propagate_span_attributes(request)
 
     # Then
-    span.set_attribute.assert_called_once_with(
-        "flagsmith.tool.name", "list_environments"
-    )
+    assert "baggage" not in request.headers