feat(MCP): Push MCP usage via OTel (#7746)

Author: emyller

api/app/settings/common.py

@@ -367,6 +367,7 @@
     "django.contrib.messages.middleware.MessageMiddleware",
     "django.middleware.clickjacking.XFrameOptionsMiddleware",
     "simple_history.middleware.HistoryRequestMiddleware",
+    "telemetry.middleware.MCPUsageLoggerMiddleware",  # Must come last!
 ]
 
 ADD_NEVER_CACHE_HEADERS = env.bool("ADD_NEVER_CACHE_HEADERS", True)

api/permissions/permission_service.py

@@ -6,6 +6,7 @@
 from environments.models import Environment
 from organisations.models import Organisation, OrganisationRole
 from projects.models import Project
+from telemetry.spans import set_span_attribute
 
 from .rbac_wrapper import (  # type: ignore[attr-defined]
     get_permitted_environments_for_master_api_key_using_roles,
@@ -25,6 +26,7 @@ def is_user_organisation_admin(
 ) -> bool:
     user_organisation = user.get_user_organisation(organisation)
     if user_organisation is not None:
+        set_span_attribute("organisation.id", user_organisation.organisation_id)
         return user_organisation.role == OrganisationRole.ADMIN.name
     return False
 

api/scripts/run-docker.sh

@@ -3,7 +3,7 @@ set -e
 
 # common environment variables
 ACCESS_LOG_FORMAT=${ACCESS_LOG_FORMAT:-'%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s "%(f)s" "%(a)s" %({origin}i)s %({access-control-allow-origin}o)s'}
-APPLICATION_LOGGERS=${APPLICATION_LOGGERS:-"app_analytics,audit,code_references,common,core,dynamodb,edge_api,environments,features,import_export,integrations,oauth2_metadata,organisations,projects,segments,task_processor,users,webhooks,workflows"}
+APPLICATION_LOGGERS=${APPLICATION_LOGGERS:-"app_analytics,audit,code_references,common,core,dynamodb,edge_api,environments,features,import_export,integrations,mcp,oauth2_metadata,organisations,projects,segments,task_processor,users,webhooks,workflows"}
 
 waitfordb() {
   if [ -z "${SKIP_WAIT_FOR_DB}" ]; then

api/telemetry/middleware.py

@@ -0,0 +1,59 @@
+from collections.abc import Callable
+
+import structlog
+from django.http.request import HttpRequest
+from django.http.response import HttpResponse
+from opentelemetry import baggage
+
+from telemetry.spans import get_span_attribute
+
+
+class MCPUsageLoggerMiddleware:
+    """Emit telemetry events for MCP usage"""
+
+    def __init__(
+        self,
+        get_response: Callable[[HttpRequest], HttpResponse],
+    ) -> None:
+        self.get_response = get_response
+
+    def __call__(self, request: HttpRequest) -> HttpResponse:
+        response = self.get_response(request)
+
+        if baggage.get_baggage("flagsmith.client.name") != "flagsmith-mcp":
+            return response
+
+        if not request.user or not request.user.is_authenticated:
+            return response
+
+        logger = structlog.get_logger("mcp")
+        event = {
+            # NOTE: The following W3C Baggage items are added by downstream processor
+            # - gen_ai.tool.name
+            # - flagsmith.mcp.client.name
+            # - flagsmith.mcp.client.version
+            "status": "error" if response.status_code >= 400 else "success",
+        }
+        if (org_id := self._get_organisation_id(request)) is not None:
+            logger.info("tool.called", organisation__id=org_id, **event)
+        else:
+            logger.warning("tool.called", organisation__id=None, **event)
+
+        return response
+
+    def _get_organisation_id(self, request: HttpRequest) -> int | None:
+        """Obtain the organisation ID from the request context."""
+        from organisations.models import Organisation
+
+        # Set by the permission layer for organisations the user belongs to
+        if isinstance(organisation_id := get_span_attribute("organisation.id"), int):
+            return organisation_id
+
+        assert request.user.is_authenticated  # NOTE: protected upstream
+        try:  # Most of the time, the user belongs to one organisation
+            return request.user.organisations.get().id
+        except (
+            Organisation.DoesNotExist,
+            Organisation.MultipleObjectsReturned,  # Don't guess
+        ):
+            return None

api/telemetry/spans.py

@@ -0,0 +1,11 @@
+from opentelemetry import trace
+from opentelemetry.util.types import AttributeValue
+
+
+def set_span_attribute(attribute: str, value: AttributeValue) -> None:
+    trace.get_current_span().set_attribute(attribute, value)
+
+
+def get_span_attribute(attribute: str) -> AttributeValue | None:
+    attributes = getattr(trace.get_current_span(), "attributes", None) or {}
+    return attributes.get(attribute)

api/tests/unit/telemetry/conftest.py

@@ -0,0 +1,22 @@
+from collections.abc import Generator
+
+import pytest
+from opentelemetry import baggage
+from opentelemetry import context as otel_context
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.trace import Span
+
+
+@pytest.fixture()
+def mcp_baggage() -> Generator[None, None, None]:
+    ctx = baggage.set_baggage("flagsmith.client.name", "flagsmith-mcp")
+    token = otel_context.attach(ctx)
+    yield
+    otel_context.detach(token)
+
+
+@pytest.fixture()
+def recording_span() -> Generator[Span, None, None]:
+    tracer = TracerProvider().get_tracer("test")
+    with tracer.start_as_current_span("test") as span:
+        yield span

api/tests/unit/telemetry/test_unit_telemetry_middleware.py

@@ -0,0 +1,168 @@
+import pytest
+from pytest_structlog import StructuredLogCapture
+from rest_framework.test import APIClient
+
+from environments.models import Environment
+from organisations.models import Organisation, OrganisationRole
+from users.models import FFAdminUser
+
+
+@pytest.mark.usefixtures("organisation")
+def test_mcp_usage_logger_middleware__no_mcp_baggage__logs_nothing(
+    staff_client: APIClient,
+    log: StructuredLogCapture,
+) -> None:
+    # Given / When
+    response = staff_client.get("/api/v1/projects/")
+
+    # Then
+    assert response.status_code == 200
+    assert log.events == []
+
+
+@pytest.mark.usefixtures("mcp_baggage", "recording_span", "organisation")
+def test_mcp_usage_logger_middleware__organisation_id_span_attribute__logs_span_organisation_id(
+    staff_client: APIClient,
+    staff_user: FFAdminUser,
+    log: StructuredLogCapture,
+) -> None:
+    # Given
+    other_organisation = Organisation.objects.create(name="Other Org")
+    staff_user.add_organisation(other_organisation, role=OrganisationRole.ADMIN)
+
+    # When
+    response = staff_client.get(
+        f"/api/v1/organisations/{other_organisation.pk}/invites/"
+    )
+
+    # Then
+    assert response.status_code == 200
+    assert log.events == [
+        {
+            "level": "info",
+            "event": "tool.called",
+            "organisation__id": other_organisation.pk,
+            "status": "success",
+        }
+    ]
+
+
+@pytest.mark.usefixtures("mcp_baggage")
+def test_mcp_usage_logger_middleware__user_with_single_organisation__logs_user_organisation_id(
+    staff_client: APIClient,
+    organisation: Organisation,
+    log: StructuredLogCapture,
+) -> None:
+    # Given / When
+    response = staff_client.get("/api/v1/projects/")
+
+    # Then
+    assert response.status_code == 200
+    assert log.events == [
+        {
+            "level": "info",
+            "event": "tool.called",
+            "organisation__id": organisation.pk,
+            "status": "success",
+        }
+    ]
+
+
+@pytest.mark.usefixtures("mcp_baggage")
+def test_mcp_usage_logger_middleware__user_without_organisations__logs_warning(
+    staff_client: APIClient,
+    staff_user: FFAdminUser,
+    log: StructuredLogCapture,
+) -> None:
+    # Given
+    assert not staff_user.organisations.exists()
+
+    # When
+    response = staff_client.get("/api/v1/projects/")
+
+    # Then
+    assert response.status_code == 200
+    assert log.events == [
+        {
+            "level": "warning",
+            "event": "tool.called",
+            "organisation__id": None,
+            "status": "success",
+        }
+    ]
+
+
+@pytest.mark.usefixtures("mcp_baggage", "organisation")
+def test_mcp_usage_logger_middleware__user_with_multiple_organisations__logs_warning(
+    staff_client: APIClient,
+    staff_user: FFAdminUser,
+    log: StructuredLogCapture,
+) -> None:
+    # Given
+    other_organisation = Organisation.objects.create(name="Other Org")
+    staff_user.add_organisation(other_organisation, role=OrganisationRole.USER)
+
+    # When
+    response = staff_client.get("/api/v1/projects/")
+
+    # Then
+    assert response.status_code == 200
+    assert log.events == [
+        {
+            "level": "warning",
+            "event": "tool.called",
+            "organisation__id": None,
+            "status": "success",
+        }
+    ]
+
+
+@pytest.mark.usefixtures("mcp_baggage")
+def test_mcp_usage_logger_middleware__unauthenticated_request__logs_nothing(
+    api_client: APIClient,
+    log: StructuredLogCapture,
+) -> None:
+    # Given / When
+    response = api_client.get("/api/v1/projects/")
+
+    # Then
+    assert response.status_code == 401
+    assert log.events == []
+
+
+@pytest.mark.usefixtures("mcp_baggage")
+def test_mcp_usage_logger_middleware__error_response__logs_error_status(
+    staff_client: APIClient,
+    organisation: Organisation,
+    log: StructuredLogCapture,
+) -> None:
+    # Given / When
+    response = staff_client.get(f"/api/v1/organisations/{organisation.pk}/invites/")
+
+    # Then
+    assert response.status_code == 403
+    assert log.events == [
+        {
+            "level": "info",
+            "event": "tool.called",
+            "organisation__id": organisation.pk,
+            "status": "error",
+        }
+    ]
+
+
+@pytest.mark.usefixtures("mcp_baggage")
+def test_mcp_usage_logger_middleware__sdk_request__logs_nothing(
+    api_client: APIClient,
+    environment: Environment,
+    log: StructuredLogCapture,
+) -> None:
+    # Given / When
+    response = api_client.get(
+        "/api/v1/flags/",
+        HTTP_X_ENVIRONMENT_KEY=environment.api_key,
+    )
+
+    # Then
+    assert response.status_code == 200
+    assert log.events == []

api/tests/unit/telemetry/test_unit_telemetry_spans.py

@@ -0,0 +1,37 @@
+import pytest
+
+from telemetry.spans import get_span_attribute, set_span_attribute
+
+
+@pytest.mark.usefixtures("recording_span")
+def test_set_span_attribute__recording_span__attribute_round_trips() -> None:
+    # Given
+    attribute = "organisation.id"
+
+    # When
+    set_span_attribute(attribute, 42)
+
+    # Then
+    assert get_span_attribute(attribute) == 42
+
+
+def test_set_span_attribute__no_recording_span__silently_ignored() -> None:
+    # Given
+    attribute = "organisation.id"
+
+    # When
+    set_span_attribute(attribute, 42)
+
+    # Then
+    assert get_span_attribute(attribute) is None
+
+
+def test_get_span_attribute__no_recording_span__returns_none() -> None:
+    # Given
+    attribute = "organisation.id"
+
+    # When
+    value = get_span_attribute(attribute)
+
+    # Then
+    assert value is None

docs/docs/deployment-self-hosting/observability/_events-catalogue.md

@@ -316,6 +316,15 @@ Attributes:
  - `project_id`
  - `retry_at`
 
+### `mcp.tool.called`
+
+Logged at `info` from:
+ - `api/telemetry/middleware.py:38`
+ - `api/telemetry/middleware.py:40`
+
+Attributes:
+ - `organisation.id`
+
 ### `platform_hub.no_analytics_database_configured`
 
 Logged at `warning` from: