refactor(MCP): Derive span attributes from baggage

Set the client identity and tool name once, as W3C Baggage, and copy
flagsmith.* entries onto spans with a span processor — replacing the
separate clientInfo span annotation. The off-the-shelf
opentelemetry-processor-baggage reads baggage from the span's parent
context, which FastMCP builds from the request _meta without baggage,
so the processor reads the current context instead.

beep boop

Author: khvn26

mcp/src/flagsmith_mcp/telemetry.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import mcp.types as mt
 from common.core.logging import setup_logging
 from common.core.otel import (
@@ -20,33 +22,55 @@
 APPLICATION_LOGGERS = ["flagsmith_mcp", "fastmcp", "mcp"]
 
 
-class ClientInfoSpanProcessor(SpanProcessor):
-    """Annotate started spans with the MCP client's self-declared identity."""
+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.
+    """
 
     def on_start(self, span: Span, parent_context: Context | None = None) -> None:
-        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)
+        for key, value in baggage.get_all().items():
+            if is_flagsmith_baggage_key(key):
+                span.set_attribute(key, str(value))
 
 
 class BaggageMiddleware(Middleware):
-    """Attach the tool name and client identity as W3C Baggage, propagated
-    to the Flagsmith API by the instrumented upstream HTTP client."""
+    """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:
-        ctx = baggage.set_baggage("flagsmith.tool.name", context.message.name)
-        if (client_info := get_client_info()) is not None:
-            ctx = baggage.set_baggage(
-                "flagsmith.client.name", client_info.name, context=ctx
-            )
-            ctx = baggage.set_baggage(
-                "flagsmith.client.version", client_info.version, context=ctx
-            )
-        token = otel_context.attach(ctx)
+        token = otel_context.attach(
+            baggage.set_baggage("flagsmith.tool.name", context.message.name)
+        )
         try:
             return await call_next(context)
         finally:
@@ -73,7 +97,7 @@ def setup_telemetry(settings: config.Settings) -> None:
             endpoint=f"{endpoint}/v1/traces",
             service_name=settings.otel_service_name,
         )
-        tracer_provider.add_span_processor(ClientInfoSpanProcessor())
+        tracer_provider.add_span_processor(FlagsmithBaggageSpanProcessor())
         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 ClientInfoSpanProcessor
+from flagsmith_mcp.telemetry import FlagsmithBaggageSpanProcessor
 
 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(ClientInfoSpanProcessor())
+    provider.add_span_processor(FlagsmithBaggageSpanProcessor())
     trace.set_tracer_provider(provider)
     return exporter
 

mcp/tests/integration/test_spans.py

@@ -29,6 +29,7 @@ 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()
@@ -47,7 +48,7 @@ async def test_spans__tool_call__upstream_request_carries_baggage(
 
     # Then the instrumented API client propagated W3C Baggage
     assert route.calls.last.request.headers["baggage"] == (
-        "flagsmith.tool.name=list_environments,"
         "flagsmith.client.name=mcp,"
-        "flagsmith.client.version=0.1.0"
+        "flagsmith.client.version=0.1.0,"
+        "flagsmith.tool.name=list_environments"
     )

mcp/tests/unit/test_telemetry.py

@@ -2,6 +2,7 @@
 
 from common.core.otel import add_otel_trace_context
 from opentelemetry import baggage
+from opentelemetry import context as otel_context
 from pytest_mock import MockerFixture
 
 from flagsmith_mcp import config, telemetry
@@ -76,7 +77,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.ClientInfoSpanProcessor)
+    assert isinstance(span_processor, telemetry.FlagsmithBaggageSpanProcessor)
     setup_logging_mock.assert_called_once_with(
         log_level="DEBUG",
         log_format="json",
@@ -88,26 +89,33 @@ def test_setup_telemetry__otlp_endpoint__exports_logs_and_traces(
     )
 
 
-def test_client_info_span_processor__outside_request_context__no_attributes(
+async def test_baggage_middleware__uninitialised_session__tool_name_baggage_only(
     mocker: MockerFixture,
 ) -> None:
-    # Given no MCP request context
-    span = mocker.Mock()
+    # 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)
 
     # When
-    telemetry.ClientInfoSpanProcessor().on_start(span)
+    await middleware.on_call_tool(context, call_next)
 
     # Then
-    span.set_attribute.assert_not_called()
+    assert seen_baggage == {"flagsmith.tool.name": "list_environments"}
+    assert baggage.get_all() == {}
 
 
-async def test_baggage_middleware__uninitialised_session__tool_name_baggage_only(
+async def test_baggage_middleware__uninitialised_session__request_baggage_untouched(
     mocker: MockerFixture,
 ) -> None:
-    # Given a tool call outside an initialised session
+    # Given a request 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:
@@ -116,8 +124,28 @@ async def record_baggage(ctx: object) -> None:
     call_next = mocker.AsyncMock(side_effect=record_baggage)
 
     # When
-    await middleware.on_call_tool(context, call_next)
+    await middleware.on_request(mocker.Mock(), call_next)
 
     # Then
-    assert seen_baggage == {"flagsmith.tool.name": "list_environments"}
-    assert baggage.get_all() == {}
+    assert seen_baggage == {}
+
+
+def test_flagsmith_baggage_span_processor__foreign_baggage__not_copied(
+    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)
+
+    # When
+    try:
+        telemetry.FlagsmithBaggageSpanProcessor().on_start(span)
+    finally:
+        otel_context.detach(token)
+
+    # Then
+    span.set_attribute.assert_called_once_with(
+        "flagsmith.tool.name", "list_environments"
+    )