diff --git a/README.md b/README.md index 9d7d444..38ccc35 100644 --- a/README.md +++ b/README.md @@ -339,7 +339,7 @@ Parameter descriptions: ### [Prometheus](https://prometheus.io/) -Prometheus integration presents a challenge because the underlying libraries for `FastAPI`, `Litestar` and `FastStream` differ significantly, making it impossible to unify them under a single interface. As a result, the Prometheus settings for `FastAPI`, `Litestar` and `FastStream` must be configured separately. +Prometheus integration presents a challenge because the underlying libraries for `FastAPI`, `Litestar` and `FastStream` differ significantly, making it impossible to unify them under a single interface. As a result, the Prometheus settings for `FastAPI`, `Litestar`, `FastStream` and `FastMCP` must be configured separately. #### FastAPI @@ -419,6 +419,43 @@ Parameters description: - `prometheus_metrics_path` - path to metrics handler. - `prometheus_middleware_cls` - Prometheus middleware for your broker. +#### FastMCP + +`FastMcpSettings` include the Prometheus settings, so the metrics route is registered by default: + +```python +from microbootstrap import FastMcpSettings + + +class YourSettings(FastMcpSettings): + service_name: str + + prometheus_metrics_path: str = "/metrics" + prometheus_register_route: bool = True + prometheus_registry: CollectorRegistry | None = None + prometheus_instrumentator_params: dict[str, typing.Any] = {} + prometheus_instrument_params: dict[str, typing.Any] = {} + prometheus_custom_labels: dict[str, typing.Any] = {} + prometheus_tool_metrics: bool = True + + ... # Other settings here +``` + +Every ASGI application returned by `http_app()` is instrumented with +[prometheus-fastapi-instrumentator](https://github.com/trallnag/prometheus-fastapi-instrumentator), the same library the FastAPI bootstrapper uses, +so requests are counted in `http_requests_total`, `http_request_duration_seconds` and friends by route template +(`handler="/mcp"`, `handler="/health/"`); requests to unknown paths are grouped into `handler="none"` and the metrics route itself is not counted. + +Parameters description: + +- `prometheus_metrics_path` - path to metrics handler. +- `prometheus_register_route` - set to `False` to serve the registry yourself. +- `prometheus_registry` - registry to expose and to register the request metrics in, the global one by default. +- `prometheus_instrumentator_params` - will be passed to `Instrumentator` during initialization (`excluded_handlers` and `registry` are set by default and can be overridden here). +- `prometheus_instrument_params` - will be passed to `Instrumentator.instrument(...)`. +- `prometheus_custom_labels` - will be added to every request and tool call metric. +- `prometheus_tool_metrics` - count tool calls in `fastmcp_tool_calls_total{tool, status}` (`status` is `success` or `error`) and `fastmcp_tool_call_duration_seconds{tool}` through a FastMCP middleware; set to `False` to turn it off. + ### [OpenTelemetry](https://opentelemetry.io/) To bootstrap OpenTelemetry, you must provide `opentelemetry_endpoint` or set `opentelemetry_log_traces` to `True`. diff --git a/microbootstrap/bootstrappers/fastmcp.py b/microbootstrap/bootstrappers/fastmcp.py index 3b7c691..356eae3 100644 --- a/microbootstrap/bootstrappers/fastmcp.py +++ b/microbootstrap/bootstrappers/fastmcp.py @@ -1,4 +1,5 @@ from __future__ import annotations +import re import typing import prometheus_client @@ -6,6 +7,8 @@ from fastmcp import FastMCP from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware from opentelemetry.util.http import ExcludeList, get_excluded_urls +from prometheus_fastapi_instrumentator import Instrumentator +from prometheus_fastapi_instrumentator import metrics as instrumentator_metrics from starlette.applications import Starlette from starlette.responses import JSONResponse, Response from starlette.routing import Match, Mount, Route @@ -18,7 +21,7 @@ from microbootstrap.instruments.prometheus_instrument import FastMcpPrometheusConfig, PrometheusInstrument from microbootstrap.instruments.pyroscope_instrument import PyroscopeInstrument from microbootstrap.instruments.sentry_instrument import SentryInstrument -from microbootstrap.middlewares.fastmcp import FastMcpLoggingMiddleware +from microbootstrap.middlewares.fastmcp import FastMcpLoggingMiddleware, FastMcpPrometheusMiddleware from microbootstrap.settings import FastMcpSettings @@ -147,6 +150,16 @@ async def health_check_handler(request: Request) -> JSONResponse: # noqa: ARG00 @FastMcpBootstrapper.use_instrument() class FastMcpPrometheusInstrument(PrometheusInstrument[FastMcpPrometheusConfig]): def bootstrap_after(self, application: FastMCP[typing.Any]) -> FastMCP[typing.Any]: # type: ignore[override] + if isinstance(application, KwargsFastMCP): + application.add_http_application_postprocessor(self.__instrument_http_app) + if self.instrument_config.prometheus_tool_metrics: + application.add_middleware( + FastMcpPrometheusMiddleware( + registry=self.instrument_config.prometheus_registry, + custom_labels=self.instrument_config.prometheus_custom_labels, + ) + ) + if not self.instrument_config.prometheus_register_route: return application @@ -165,6 +178,29 @@ async def metrics_handler(request: Request) -> Response: # noqa: ARG001 return application + _default_metrics: typing.Callable[..., None] | None = None + + def __instrument_http_app(self, http_application: StarletteT) -> StarletteT: + # Same instrumentator as in the FastAPI bootstrapper: requests are counted by route template + # (`handler="/mcp"`, `handler="/health/"`), unknown paths are grouped into `handler="none"`. + registry: typing.Final = self.instrument_config.prometheus_registry or prometheus_client.REGISTRY + instrumentator_params: typing.Final = { + "excluded_handlers": [f"^{re.escape(self.instrument_config.prometheus_metrics_path)}$"], + "registry": registry, + **self.instrument_config.prometheus_instrumentator_params, + } + # `metrics.default()` registers its collectors on creation and returns None on a duplicate, so a second + # `http_app()` would get an instrumentator without metrics. The closure is built once and shared instead. + if self._default_metrics is None: + self._default_metrics = instrumentator_metrics.default( + registry=registry, + custom_labels=self.instrument_config.prometheus_custom_labels, + ) + Instrumentator(**instrumentator_params).add(self._default_metrics).instrument( + http_application, **self.instrument_config.prometheus_instrument_params + ) + return http_application + @classmethod def get_config_type(cls) -> type[FastMcpPrometheusConfig]: return FastMcpPrometheusConfig diff --git a/microbootstrap/instruments/prometheus_instrument.py b/microbootstrap/instruments/prometheus_instrument.py index 7d1f290..0a7332e 100644 --- a/microbootstrap/instruments/prometheus_instrument.py +++ b/microbootstrap/instruments/prometheus_instrument.py @@ -35,6 +35,10 @@ class FastApiPrometheusConfig(BasePrometheusConfig): class FastMcpPrometheusConfig(BasePrometheusConfig): prometheus_registry: typing.Any | None = None prometheus_register_route: bool = True + prometheus_instrumentator_params: dict[str, typing.Any] = pydantic.Field(default_factory=dict) + prometheus_instrument_params: dict[str, typing.Any] = pydantic.Field(default_factory=dict) + prometheus_custom_labels: dict[str, typing.Any] = pydantic.Field(default_factory=dict) + prometheus_tool_metrics: bool = True @typing.runtime_checkable diff --git a/microbootstrap/middlewares/fastmcp.py b/microbootstrap/middlewares/fastmcp.py index 488aa40..af024fd 100644 --- a/microbootstrap/middlewares/fastmcp.py +++ b/microbootstrap/middlewares/fastmcp.py @@ -2,12 +2,15 @@ import time import typing +import prometheus_client import structlog from fastmcp.server.middleware import Middleware, MiddlewareContext if typing.TYPE_CHECKING: + import mcp.types as mcp_types from fastmcp.server.middleware import CallNext + from fastmcp.tools.base import ToolResult fastmcp_access_logger: typing.Final = structlog.get_logger("mcp.access") @@ -44,3 +47,50 @@ async def on_message( duration=time.perf_counter_ns() - start_time, ) return result + + +TOOL_CALL_STATUS_SUCCESS: typing.Final = "success" +TOOL_CALL_STATUS_ERROR: typing.Final = "error" + + +class FastMcpPrometheusMiddleware(Middleware): + def __init__( + self, + registry: prometheus_client.CollectorRegistry | None = None, + custom_labels: dict[str, str] | None = None, + duration_buckets: typing.Sequence[float] = (0.01, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, float("inf")), + ) -> None: + self.custom_labels: typing.Final = custom_labels or {} + self.tool_calls_total: typing.Final = prometheus_client.Counter( + name="fastmcp_tool_calls_total", + documentation="Number of MCP tool calls by tool and outcome.", + labelnames=["tool", "status", *self.custom_labels], + registry=registry or prometheus_client.REGISTRY, + ) + self.tool_call_duration_seconds: typing.Final = prometheus_client.Histogram( + name="fastmcp_tool_call_duration_seconds", + documentation="Duration of MCP tool calls by tool.", + labelnames=["tool", *self.custom_labels], + buckets=duration_buckets, + registry=registry or prometheus_client.REGISTRY, + ) + + async def on_call_tool( + self, + context: MiddlewareContext[mcp_types.CallToolRequestParams], + call_next: CallNext[mcp_types.CallToolRequestParams, ToolResult], + ) -> ToolResult: + tool_name: typing.Final = context.message.name + start_time: typing.Final = time.perf_counter() + try: + result: typing.Final = await call_next(context) + except Exception: + self.tool_calls_total.labels(tool=tool_name, status=TOOL_CALL_STATUS_ERROR, **self.custom_labels).inc() + raise + else: + self.tool_calls_total.labels(tool=tool_name, status=TOOL_CALL_STATUS_SUCCESS, **self.custom_labels).inc() + return result + finally: + self.tool_call_duration_seconds.labels(tool=tool_name, **self.custom_labels).observe( + time.perf_counter() - start_time + ) diff --git a/pyproject.toml b/pyproject.toml index b29a08d..c156a7c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,7 +75,7 @@ litestar = [ ] granian = ["granian[reload]>=1"] faststream = ["faststream>=0.6.7,<0.8", "prometheus-client>=0.20"] -fastmcp = ["fastmcp>=2,<4", "prometheus-client>=0.20"] +fastmcp = ["fastmcp>=2,<4", "prometheus-client>=0.20", "prometheus-fastapi-instrumentator>=7.1"] [dependency-groups] dev = [ diff --git a/tests/bootstrappers/test_fastmcp.py b/tests/bootstrappers/test_fastmcp.py index afbf454..969c12f 100644 --- a/tests/bootstrappers/test_fastmcp.py +++ b/tests/bootstrappers/test_fastmcp.py @@ -4,7 +4,8 @@ import prometheus_client import pytest -from fastmcp import FastMCP +from fastmcp import Client, FastMCP +from fastmcp.exceptions import ToolError from opentelemetry.instrumentation._semconv import ( OTEL_SEMCONV_STABILITY_OPT_IN, _OpenTelemetrySemanticConventionStability, @@ -30,7 +31,7 @@ from microbootstrap.instruments.logging_instrument import LoggingConfig from microbootstrap.instruments.opentelemetry_instrument import OpentelemetryConfig, OpenTelemetryInstrumentor from microbootstrap.instruments.prometheus_instrument import FastMcpPrometheusConfig -from microbootstrap.middlewares.fastmcp import FastMcpLoggingMiddleware +from microbootstrap.middlewares.fastmcp import FastMcpLoggingMiddleware, FastMcpPrometheusMiddleware from microbootstrap.settings import FastMcpSettings @@ -152,6 +153,153 @@ def test_fastmcp_prometheus() -> None: assert b"fastmcp_test_requests_total 1.0" in response.content +def test_fastmcp_prometheus_counts_http_requests_by_route_template() -> None: + test_metrics_path: typing.Final = "/test-metrics" + test_health_path: typing.Final = "/test-health/" + metrics_registry: typing.Final = prometheus_client.CollectorRegistry() + application: typing.Final = ( + FastMcpBootstrapper(FastMcpSettings()) + .configure_instrument( + FastMcpPrometheusConfig( + prometheus_metrics_path=test_metrics_path, + prometheus_registry=metrics_registry, + prometheus_custom_labels={"team": "platform"}, + ), + ) + .configure_instrument(HealthChecksConfig(health_checks_path=test_health_path)) + .bootstrap() + ) + + with TestClient(application.http_app(path="/mcp")) as client: + for _ in range(2): + assert client.get(test_health_path).status_code == status.HTTP_200_OK + client.get("/unknown") + client.post("/mcp", json={}) + assert client.get(test_metrics_path).status_code == status.HTTP_200_OK + + def requests_total(method: str, status_group: str, handler: str) -> float | None: + return metrics_registry.get_sample_value( + "http_requests_total", + {"method": method, "status": status_group, "handler": handler, "team": "platform"}, + ) + + assert requests_total("GET", "2xx", test_health_path) == 2 # noqa: PLR2004 + assert requests_total("GET", "4xx", "none") == 1 + assert requests_total("POST", "4xx", "/mcp") == 1 + assert requests_total("GET", "2xx", test_metrics_path) is None + + +def test_fastmcp_prometheus_instrumentator_params_are_passed() -> None: + metrics_registry: typing.Final = prometheus_client.CollectorRegistry() + application: typing.Final = ( + FastMcpBootstrapper(FastMcpSettings()) + .configure_instrument( + FastMcpPrometheusConfig( + prometheus_registry=metrics_registry, + prometheus_instrumentator_params={"should_group_status_codes": False, "excluded_handlers": []}, + ), + ) + .bootstrap() + ) + + with TestClient(application.http_app()) as client: + assert client.get("/metrics").status_code == status.HTTP_200_OK + + # excluded_handlers is overridden, so the metrics route itself is counted, with the exact status code + assert ( + metrics_registry.get_sample_value( + "http_requests_total", + {"method": "GET", "status": "200", "handler": "/metrics"}, + ) + == 1 + ) + + +def test_fastmcp_prometheus_counts_requests_of_every_http_application() -> None: + test_health_path: typing.Final = "/test-health/" + metrics_registry: typing.Final = prometheus_client.CollectorRegistry() + application: typing.Final = ( + FastMcpBootstrapper(FastMcpSettings()) + .configure_instrument(FastMcpPrometheusConfig(prometheus_registry=metrics_registry)) + .configure_instrument(HealthChecksConfig(health_checks_path=test_health_path)) + .bootstrap() + ) + + # A second http_app() (another transport, or a test building the app twice) must be counted as well: + # the instrumentator's default metrics register once and are silently skipped on a duplicate + for _ in range(2): + with TestClient(application.http_app()) as client: + assert client.get(test_health_path).status_code == status.HTTP_200_OK + + assert ( + metrics_registry.get_sample_value( + "http_requests_total", + {"method": "GET", "status": "2xx", "handler": test_health_path}, + ) + == 2 # noqa: PLR2004 + ) + + +async def test_fastmcp_prometheus_counts_tool_calls() -> None: + metrics_registry: typing.Final = prometheus_client.CollectorRegistry() + application: typing.Final = ( + FastMcpBootstrapper(FastMcpSettings()) + .configure_instrument( + FastMcpPrometheusConfig( + prometheus_registry=metrics_registry, prometheus_custom_labels={"team": "platform"} + ), + ) + .bootstrap() + ) + + @application.tool + def echo(text: str) -> str: + return text + + @application.tool + def failing() -> str: + msg: typing.Final = "boom" + raise ValueError(msg) + + async with Client(application) as client: + for _ in range(2): + await client.call_tool("echo", {"text": "hi"}) + with pytest.raises(ToolError): + await client.call_tool("failing", {}) + + def tool_calls_total(tool: str, call_status: str) -> float | None: + return metrics_registry.get_sample_value( + "fastmcp_tool_calls_total", + {"tool": tool, "status": call_status, "team": "platform"}, + ) + + assert tool_calls_total("echo", "success") == 2 # noqa: PLR2004 + assert tool_calls_total("echo", "error") is None + assert tool_calls_total("failing", "error") == 1 + assert ( + metrics_registry.get_sample_value( + "fastmcp_tool_call_duration_seconds_count", + {"tool": "failing", "team": "platform"}, + ) + == 1 + ) + + +def test_fastmcp_prometheus_tool_metrics_can_be_disabled() -> None: + application: typing.Final = ( + FastMcpBootstrapper(FastMcpSettings()) + .configure_instrument( + FastMcpPrometheusConfig( + prometheus_registry=prometheus_client.CollectorRegistry(), + prometheus_tool_metrics=False, + ), + ) + .bootstrap() + ) + + assert not any(isinstance(middleware, FastMcpPrometheusMiddleware) for middleware in application.middleware) + + def test_fastmcp_prometheus_route_can_be_disabled() -> None: test_metrics_path: typing.Final = "/test-metrics" application: typing.Final = ( @@ -230,7 +378,10 @@ def test_fastmcp_opentelemetry_is_not_ready_without_settings() -> None: application: typing.Final = FastMcpBootstrapper(FastMcpSettings(service_debug=False)).bootstrap() assert isinstance(application, KwargsFastMCP) - assert application.http_application_postprocessors == [] + assert not any( + isinstance(getattr(postprocessor, "__self__", None), FastMcpOpentelemetryInstrument) + for postprocessor in application.http_application_postprocessors + ) assert count_opentelemetry_middlewares(application.http_app()) == 0