Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions api/app/settings/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions api/permissions/permission_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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

Expand Down
2 changes: 1 addition & 1 deletion api/scripts/run-docker.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
59 changes: 59 additions & 0 deletions api/telemetry/middleware.py
Original file line number Diff line number Diff line change
@@ -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
11 changes: 11 additions & 0 deletions api/telemetry/spans.py
Original file line number Diff line number Diff line change
@@ -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)
22 changes: 22 additions & 0 deletions api/tests/unit/telemetry/conftest.py
Original file line number Diff line number Diff line change
@@ -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
168 changes: 168 additions & 0 deletions api/tests/unit/telemetry/test_unit_telemetry_middleware.py
Original file line number Diff line number Diff line change
@@ -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 == []
37 changes: 37 additions & 0 deletions api/tests/unit/telemetry/test_unit_telemetry_spans.py
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading