feat(MCP): Expose Prometheus metrics (#7705)

Author: khvn26

mcp/pyproject.toml

@@ -4,24 +4,27 @@ version = "0.1.0"
 description = "An MCP server connecting to the Flagsmith API."
 authors = [{ name = "Flagsmith", email = "support@flagsmith.com" }]
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
-    "fastmcp>=3.3.1,<4.0.0",           # Base MCP functionality
-    "pydantic-settings>=2.0.0,<3.0.0", # Environment-driven configuration
+    "fastmcp>=3.3.1,<4.0.0",            # Base MCP functionality
+    "prometheus-client>=0.21.0,<1.0.0", # Export Prometheus metrics
+    "pydantic-settings>=2.0.0,<3.0.0",  # Environment-driven configuration
 ]
 
 [project.scripts]
 flagsmith-mcp = "flagsmith_mcp.server:run"
 
 [dependency-groups]
 dev = [
-    "mypy>=2.1.0,<3.0.0",             # Static type checking
-    "openapi-pydantic>=0.5.0,<1.0.0", # Build OpenAPI specs as fixtures
-    "pytest>=9.0.3,<10.0.0",          # Run tests
-    "pytest-asyncio>=1.3.0,<2.0.0",   # Run asynchronous tests
-    "pytest-cov>=7.0.0,<8.0.0",       # Measure test coverage
-    "respx>=0.22,<1.0",               # Mock HTTP interactions
-    "ruff>=0.15.12,<0.16.0",          # Lint and format
+    "flagsmith-common[test-tools]>=3.9.1,<4.0.0", # Shared test fixtures
+    "mypy>=2.1.0,<3.0.0",                         # Static type checking
+    "openapi-pydantic>=0.5.0,<1.0.0",             # Build OpenAPI specs as fixtures
+    "pytest>=9.0.3,<10.0.0",                      # Run tests
+    "pytest-asyncio>=1.3.0,<2.0.0",               # Run asynchronous tests
+    "pytest-cov>=7.0.0,<8.0.0",                   # Measure test coverage
+    "pytest-mock>=3.15.1,<4.0.0",                 # Mock via fixtures
+    "respx>=0.22,<1.0",                           # Mock HTTP interactions
+    "ruff>=0.15.12,<0.16.0",                      # Lint and format
 ]
 
 [build-system]
@@ -37,7 +40,7 @@ branch = true
 
 [tool.ruff]
 line-length = 88
-target-version = "py310"
+target-version = "py311"
 
 [tool.ruff.lint]
 # Establish parity with flake8 + isort

mcp/src/flagsmith_mcp/config.py

@@ -21,6 +21,10 @@ class Settings(BaseSettings):
         default="http",
     )
     """MCP transport to use."""
+    metrics_port: int | None = Field(
+        default=None,
+    )
+    """Serve Prometheus metrics on this port. Disabled when unset."""
     mcp_server_url: str = Field(
         default="http://127.0.0.1:8000",
     )

mcp/src/flagsmith_mcp/metrics.py

@@ -0,0 +1,89 @@
+import time
+from collections.abc import Sequence
+
+import mcp.types as mt
+import pydantic_core
+from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
+from fastmcp.tools.base import Tool, ToolResult
+from prometheus_client import Gauge, Histogram
+
+flagsmith_mcp_tool_call_duration_seconds = Histogram(
+    "flagsmith_mcp_tool_call_duration_seconds",
+    "Time spent serving an MCP tool call, including the upstream "
+    "Flagsmith API request.",
+    labelnames=["tool", "status"],
+)
+flagsmith_mcp_tool_result_bytes = Histogram(
+    "flagsmith_mcp_tool_result_bytes",
+    "Serialised size of a tool call result's text content and "
+    "structuredContent. Helps approximate the call's token cost: clients "
+    "render either or both into the agent's context.",
+    labelnames=["tool", "content"],
+    buckets=(256, 1024, 4096, 16384, 65536, 262144, 1048576, float("inf")),
+)
+flagsmith_mcp_tool_catalogue_bytes = Gauge(
+    "flagsmith_mcp_tool_catalogue_bytes",
+    "Serialised size of the tool catalogue returned by tools/list. A proxy "
+    "for the token cost every MCP session pays before any tool is called.",
+)
+
+
+class PrometheusMiddleware(Middleware):
+    """Record Prometheus metrics for MCP tool calls."""
+
+    async def on_call_tool(
+        self,
+        context: MiddlewareContext[mt.CallToolRequestParams],
+        call_next: CallNext[mt.CallToolRequestParams, ToolResult],
+    ) -> ToolResult:
+        tool = context.message.name
+        start = time.perf_counter()
+        try:
+            result = await call_next(context)
+        except Exception:
+            flagsmith_mcp_tool_call_duration_seconds.labels(
+                tool=tool, status="error"
+            ).observe(time.perf_counter() - start)
+            raise
+        flagsmith_mcp_tool_call_duration_seconds.labels(
+            tool=tool, status="success"
+        ).observe(time.perf_counter() - start)
+        # Text blocks already hold the serialised payload; structuredContent
+        # costs one compact dump, matching its wire encoding.
+        unstructured_bytes = sum(
+            len(block.text.encode())
+            for block in result.content
+            if isinstance(block, mt.TextContent)
+        )
+        structured_bytes = (
+            len(pydantic_core.to_json(result.structured_content, fallback=str))
+            if result.structured_content is not None
+            else 0
+        )
+        for content, size in (
+            ("unstructured", unstructured_bytes),
+            ("structured", structured_bytes),
+        ):
+            flagsmith_mcp_tool_result_bytes.labels(
+                tool=tool,
+                content=content,
+            ).observe(size)
+        return result
+
+    async def on_list_tools(
+        self,
+        context: MiddlewareContext[mt.ListToolsRequest],
+        call_next: CallNext[mt.ListToolsRequest, Sequence[Tool]],
+    ) -> Sequence[Tool]:
+        tools = await call_next(context)
+        flagsmith_mcp_tool_catalogue_bytes.set(
+            len(
+                pydantic_core.to_json(
+                    [
+                        tool.to_mcp_tool().model_dump(exclude_none=True, by_alias=True)
+                        for tool in tools
+                    ]
+                )
+            )
+        )
+        return tools

mcp/src/flagsmith_mcp/server.py

@@ -6,11 +6,13 @@
 from fastmcp.utilities.components import FastMCPComponent
 from fastmcp.utilities.openapi.models import HttpMethod, HTTPRoute
 from mcp.types import ToolAnnotations
+from prometheus_client import start_http_server
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse
 
 from flagsmith_mcp import config, constants
 from flagsmith_mcp.auth import FlagsmithAuth
+from flagsmith_mcp.metrics import PrometheusMiddleware
 from flagsmith_mcp.oauth import FlagsmithResourceAuth
 
 ROUTE_MAPS = [
@@ -66,6 +68,8 @@ def create_server(settings: config.Settings) -> FastMCP[None]:
         auth=auth,
     )
 
+    server.add_middleware(PrometheusMiddleware())
+
     @server.custom_route("/health", methods=["GET"])
     async def health(request: Request) -> PlainTextResponse:
         return PlainTextResponse("OK")
@@ -75,4 +79,7 @@ async def health(request: Request) -> PlainTextResponse:
 
 def run() -> None:
     settings = config.Settings()
-    create_server(settings).run(transport=settings.transport)
+    server = create_server(settings)
+    if settings.metrics_port is not None:
+        start_http_server(settings.metrics_port)
+    server.run(transport=settings.transport)

mcp/tests/integration/test_metrics.py

@@ -0,0 +1,81 @@
+import pytest
+from common.test_tools import AssertMetricFixture
+from fastmcp import Client
+from fastmcp.client.transports import FastMCPTransport
+from fastmcp.exceptions import ToolError
+from prometheus_client import REGISTRY
+from respx import MockRouter
+
+
+async def test_metrics__successful_tool_call__records_duration_and_result_size(
+    client: Client[FastMCPTransport],
+    respx_mock: MockRouter,
+    assert_metric: AssertMetricFixture,
+) -> None:
+    # Given
+    respx_mock.get("https://api.flagsmith.com/environments/").respond(
+        json={"results": []}
+    )
+
+    # When
+    await client.call_tool("list_environments", {})
+
+    # Then
+    assert_metric(
+        name="flagsmith_mcp_tool_call_duration_seconds_count",
+        labels={"tool": "list_environments", "status": "success"},
+        value=1,
+    )
+    for content in ("unstructured", "structured"):
+        assert_metric(
+            name="flagsmith_mcp_tool_result_bytes_count",
+            labels={"tool": "list_environments", "content": content},
+            value=1,
+        )
+        content_sum = REGISTRY.get_sample_value(
+            "flagsmith_mcp_tool_result_bytes_sum",
+            {"tool": "list_environments", "content": content},
+        )
+        assert content_sum is not None
+        assert content_sum > 0
+
+
+async def test_metrics__failing_tool_call__records_error_duration_only(
+    client: Client[FastMCPTransport],
+    respx_mock: MockRouter,
+    assert_metric: AssertMetricFixture,
+) -> None:
+    # Given
+    respx_mock.get("https://api.flagsmith.com/environments/").respond(status_code=502)
+
+    # When
+    with pytest.raises(ToolError):
+        await client.call_tool("list_environments", {})
+
+    # Then
+    assert_metric(
+        name="flagsmith_mcp_tool_call_duration_seconds_count",
+        labels={"tool": "list_environments", "status": "error"},
+        value=1,
+    )
+    assert (
+        REGISTRY.get_sample_value(
+            "flagsmith_mcp_tool_result_bytes_count",
+            {"tool": "list_environments", "content": "unstructured"},
+        )
+        is None
+    )
+
+
+async def test_metrics__tools_list__records_catalogue_size(
+    client: Client[FastMCPTransport],
+) -> None:
+    # Given the server started via the client fixture
+    # When
+    tools = await client.list_tools()
+
+    # Then a catalogue of two tools weighs at least a name and a schema each
+    assert tools
+    catalogue_bytes = REGISTRY.get_sample_value("flagsmith_mcp_tool_catalogue_bytes")
+    assert catalogue_bytes is not None
+    assert catalogue_bytes > len(tools) * 50

mcp/tests/unit/test_server.py

@@ -1,9 +1,10 @@
-from typing import Any
+import os
 
 import openapi_pydantic as openapi
 import pytest
 from fastmcp import Client
 from mcp.types import ToolAnnotations
+from pytest_mock import MockerFixture
 from respx import MockRouter
 
 from flagsmith_mcp import config, constants, server
@@ -124,21 +125,52 @@ async def test_create_server__untagged_route__excluded_from_tools(
 
 
 def test_run__configured_transport__runs_server_with_it(
-    monkeypatch: pytest.MonkeyPatch,
+    mocker: MockerFixture,
 ) -> None:
     # Given
-    calls: dict[str, Any] = {}
+    mocker.patch.dict(
+        os.environ,
+        {"TRANSPORT": "stdio", "FLAGSMITH_API_TOKEN": "ser.secret"},
+        clear=True,
+    )
+    create_server_mock = mocker.patch.object(server, "create_server", autospec=True)
+
+    # When
+    server.run()
+
+    # Then
+    create_server_mock.return_value.run.assert_called_once_with(transport="stdio")
+
 
-    class FakeServer:
-        def run(self, transport: str) -> None:
-            calls["transport"] = transport
+def test_run__metrics_port_unset__metrics_server_not_started(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch.dict(os.environ, {}, clear=True)
+    mocker.patch.object(server, "create_server", autospec=True)
+    start_http_server_mock = mocker.patch.object(
+        server, "start_http_server", autospec=True
+    )
 
-    monkeypatch.setattr(server, "create_server", lambda settings: FakeServer())
-    monkeypatch.setenv("TRANSPORT", "stdio")
-    monkeypatch.setenv("FLAGSMITH_API_TOKEN", "ser.secret")
+    # When
+    server.run()
+
+    # Then
+    start_http_server_mock.assert_not_called()
+
+
+def test_run__metrics_port_set__metrics_server_started_with_it(
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    mocker.patch.dict(os.environ, {"METRICS_PORT": "9464"}, clear=True)
+    mocker.patch.object(server, "create_server", autospec=True)
+    start_http_server_mock = mocker.patch.object(
+        server, "start_http_server", autospec=True
+    )
 
     # When
     server.run()
 
     # Then
-    assert calls == {"transport": "stdio"}
+    start_http_server_mock.assert_called_once_with(9464)