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
25 changes: 14 additions & 11 deletions mcp/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Comment thread
emyller marked this conversation as resolved.
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]
Expand All @@ -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
Expand Down
4 changes: 4 additions & 0 deletions mcp/src/flagsmith_mcp/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)
Expand Down
89 changes: 89 additions & 0 deletions mcp/src/flagsmith_mcp/metrics.py
Original file line number Diff line number Diff line change
@@ -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
9 changes: 8 additions & 1 deletion mcp/src/flagsmith_mcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -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")
Expand All @@ -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)
81 changes: 81 additions & 0 deletions mcp/tests/integration/test_metrics.py
Original file line number Diff line number Diff line change
@@ -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
Comment thread
emyller marked this conversation as resolved.
52 changes: 42 additions & 10 deletions mcp/tests/unit/test_server.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
Loading
Loading