diff --git a/api/app/settings/common.py b/api/app/settings/common.py index ab4e9c8ecec3..d52655168c65 100644 --- a/api/app/settings/common.py +++ b/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) diff --git a/api/permissions/permission_service.py b/api/permissions/permission_service.py index 1d37ed15d387..07b25b4124c4 100644 --- a/api/permissions/permission_service.py +++ b/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 diff --git a/api/scripts/run-docker.sh b/api/scripts/run-docker.sh index ffe8b2207487..35cc39d63496 100755 --- a/api/scripts/run-docker.sh +++ b/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 diff --git a/api/telemetry/middleware.py b/api/telemetry/middleware.py new file mode 100644 index 000000000000..7754a686f523 --- /dev/null +++ b/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 diff --git a/api/telemetry/spans.py b/api/telemetry/spans.py new file mode 100644 index 000000000000..c4e60e1f4359 --- /dev/null +++ b/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) diff --git a/api/tests/unit/telemetry/conftest.py b/api/tests/unit/telemetry/conftest.py new file mode 100644 index 000000000000..2d2eb39f5bf2 --- /dev/null +++ b/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 diff --git a/api/tests/unit/telemetry/test_unit_telemetry_middleware.py b/api/tests/unit/telemetry/test_unit_telemetry_middleware.py new file mode 100644 index 000000000000..2ce11324edc0 --- /dev/null +++ b/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 == [] diff --git a/api/tests/unit/telemetry/test_unit_telemetry_spans.py b/api/tests/unit/telemetry/test_unit_telemetry_spans.py new file mode 100644 index 000000000000..82b9d19037d2 --- /dev/null +++ b/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 diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index 38400b6582ba..e4615245738e 100644 --- a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md +++ b/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: