feat(MCP): Propagate client identity and tool name to the API

Instrument the Flagsmith API httpx client (instance-only, via
opentelemetry-instrumentation-httpx) so every upstream call emits a
client span and carries W3C trace context and baggage. A middleware
sets flagsmith.tool.name and flagsmith.client.name/version as baggage
per tool call; the API's existing OTel pipeline extracts baggage and
copies it onto its own events, giving per-client, per-tool visibility
on both sides of the boundary.

beep boop

Author: khvn26

mcp/pyproject.toml

@@ -8,6 +8,7 @@ requires-python = ">=3.11"
 dependencies = [
     "fastmcp>=3.3.1,<4.0.0",                    # Base MCP functionality
     "flagsmith-common[otel]>=3.10.0,<4.0.0",    # Logging and OTel export
+    "opentelemetry-instrumentation-httpx>=0.46b0,<1.0.0", # Trace upstream API calls
     "prometheus-client>=0.21.0,<1.0.0",         # Export Prometheus metrics
     "pydantic-settings>=2.0.0,<3.0.0",          # Environment-driven configuration
 ]

mcp/src/flagsmith_mcp/server.py

@@ -6,6 +6,7 @@
 from fastmcp.utilities.components import FastMCPComponent
 from fastmcp.utilities.openapi.models import HttpMethod, HTTPRoute
 from mcp.types import ToolAnnotations
+from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
 from prometheus_client import start_http_server
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse
@@ -15,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 setup_telemetry
+from flagsmith_mcp.telemetry import BaggageMiddleware, setup_telemetry
 
 ROUTE_MAPS = [
     RouteMap(tags={"mcp"}, mcp_type=MCPType.TOOL),
@@ -57,12 +58,16 @@ def create_server(settings: config.Settings) -> FastMCP[None]:
             resource_url=settings.mcp_server_url,
             authorization_server=settings.flagsmith_api_url,
         )
+    api_client = httpx.AsyncClient(
+        base_url=settings.flagsmith_api_url,
+        auth=FlagsmithAuth(settings.flagsmith_api_token),
+    )
+    # Instrument only the Flagsmith API client: emit a span per upstream
+    # call and propagate W3C trace context and baggage to the API.
+    HTTPXClientInstrumentor().instrument_client(api_client)
     server = FastMCP.from_openapi(
         openapi_spec=_fetch_spec(),
-        client=httpx.AsyncClient(
-            base_url=settings.flagsmith_api_url,
-            auth=FlagsmithAuth(settings.flagsmith_api_token),
-        ),
+        client=api_client,
         name="Flagsmith",
         route_maps=ROUTE_MAPS,
         mcp_component_fn=_customise,
@@ -72,6 +77,7 @@ 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,11 +1,15 @@
+import mcp.types as mt
 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 opentelemetry import trace
+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.context import Context
 from opentelemetry.sdk.trace import Span, SpanProcessor
 from structlog.typing import Processor
@@ -25,6 +29,30 @@ def on_start(self, span: Span, parent_context: Context | None = None) -> None:
             span.set_attribute("flagsmith.client.version", client_info.version)
 
 
+class BaggageMiddleware(Middleware):
+    """Attach the tool name and client identity as W3C Baggage, propagated
+    to the Flagsmith API by the instrumented upstream HTTP client."""
+
+    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)
+        try:
+            return await call_next(context)
+        finally:
+            otel_context.detach(token)
+
+
 def setup_telemetry(settings: config.Settings) -> None:
     """Set up logging, exporting structlog events and traces to OpenTelemetry
     when an OTLP endpoint is configured."""

mcp/tests/integration/test_spans.py

@@ -32,3 +32,22 @@ async def test_spans__tool_call__annotated_with_client_identity(
         "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(
+    client: Client[FastMCPTransport],
+    respx_mock: MockRouter,
+) -> None:
+    # Given
+    route = respx_mock.get("https://api.flagsmith.com/environments/")
+    route.respond(json={"results": []})
+
+    # When
+    await client.call_tool("list_environments", {})
+
+    # 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"
+    )

mcp/tests/unit/test_telemetry.py

@@ -1,6 +1,7 @@
 import os
 
 from common.core.otel import add_otel_trace_context
+from opentelemetry import baggage
 from pytest_mock import MockerFixture
 
 from flagsmith_mcp import config, telemetry
@@ -98,3 +99,25 @@ def test_client_info_span_processor__outside_request_context__no_attributes(
 
     # Then
     span.set_attribute.assert_not_called()
+
+
+async def test_baggage_middleware__uninitialised_session__tool_name_baggage_only(
+    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)
+
+    # When
+    await middleware.on_call_tool(context, call_next)
+
+    # Then
+    assert seen_baggage == {"flagsmith.tool.name": "list_environments"}
+    assert baggage.get_all() == {}