From 9c5e92e219db413fc11a9392986573d3153c6e88 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Tue, 10 Nov 2026 12:58:53 +0300 Subject: [PATCH] Use the configured Prometheus registry in FastMCP Closes #451 --- docs/migration/to-2.x.md | 8 ++ .../bootstrappers/fastmcp_bootstrapper.py | 89 +++++++++++++++--- tests/test_fastmcp_bootstrap.py | 94 +++++++++++++++++++ 3 files changed, 179 insertions(+), 12 deletions(-) diff --git a/docs/migration/to-2.x.md b/docs/migration/to-2.x.md index f4ce04f..67c535e 100644 --- a/docs/migration/to-2.x.md +++ b/docs/migration/to-2.x.md @@ -331,6 +331,14 @@ headers, including `Content-Type`. Add `Authorization` explicitly when needed ## Prometheus +### FastMCP custom registry + +When `prometheus_instrumentator_additional_params["registry"]` supplies a custom +registry, FastMCP's Prometheus HTTP and tool metrics use that registry. The +metrics endpoint exposes its collectors. Register any application collectors you +want to expose in that same registry. FastMCP uses the global registry by default +([#451](https://github.com/modern-python/lite-bootstrap/issues/451)). + ### Unknown FastMCP tools Calls to unresolved tool names no longer create `fastmcp_tool_calls_total` or diff --git a/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py b/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py index 57c89e6..edbf86b 100644 --- a/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py @@ -2,8 +2,10 @@ import dataclasses import functools import inspect +import threading import time import typing +from _thread import LockType from collections.abc import AsyncGenerator, Mapping from lite_bootstrap import import_checker @@ -27,8 +29,10 @@ from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.server.providers.addressing import parse_hashed_backend_name from starlette import status + from starlette.middleware import Middleware as StarletteMiddleware from starlette.requests import Request from starlette.responses import JSONResponse, Response + from starlette.types import ASGIApp if import_checker.is_fastmcp_opentelemetry_installed: from opentelemetry.metrics import get_meter_provider @@ -53,9 +57,23 @@ if import_checker.is_prometheus_fastapi_instrumentator_installed: from prometheus_fastapi_instrumentator import Instrumentator from prometheus_fastapi_instrumentator import metrics as instrumentator_metrics + from prometheus_fastapi_instrumentator.middleware import PrometheusInstrumentatorMiddleware _DEFAULT_METRICS_PARAMETERS: typing.Final = frozenset(inspect.signature(instrumentator_metrics.default).parameters) + class _FastMcpPrometheusHttpMiddleware(PrometheusInstrumentatorMiddleware): + def __init__( + self, + app: "ASGIApp", + *, + inprogress_gauge: "prometheus_client.Gauge", + **kwargs: typing.Any, # noqa: ANN401 + ) -> None: + kwargs["should_instrument_requests_inprogress"] = False + super().__init__(app, **kwargs) + self.inprogress = inprogress_gauge + self.should_instrument_requests_inprogress = True + _TOOL_CALL_STATUS_SUCCESS: typing.Final = "success" _TOOL_CALL_STATUS_ERROR: typing.Final = "error" @@ -94,9 +112,11 @@ def restore() -> None: return restore -def _build_tool_call_metrics() -> tuple["prometheus_client.Counter", "prometheus_client.Histogram"]: - # The global registry rejects a second collector of the same name, so a later bootstrap reuses the first - registered_collectors: typing.Final = prometheus_client.REGISTRY._names_to_collectors # noqa: SLF001 +def _build_tool_call_metrics( + registry: "prometheus_client.CollectorRegistry", +) -> tuple["prometheus_client.Counter", "prometheus_client.Histogram"]: + # The registry rejects a second collector of the same name, so a later bootstrap reuses the first + registered_collectors: typing.Final = registry._names_to_collectors # noqa: SLF001 if "fastmcp_tool_calls_total" in registered_collectors: return ( typing.cast("prometheus_client.Counter", registered_collectors["fastmcp_tool_calls_total"]), @@ -104,10 +124,13 @@ def _build_tool_call_metrics() -> tuple["prometheus_client.Counter", "prometheus ) return ( prometheus_client.Counter( - "fastmcp_tool_calls_total", "Number of MCP tool calls by tool and outcome.", ["tool", "status"] + "fastmcp_tool_calls_total", + "Number of MCP tool calls by tool and outcome.", + ["tool", "status"], + registry=registry, ), prometheus_client.Histogram( - "fastmcp_tool_call_duration_seconds", "Duration of MCP tool calls by tool.", ["tool"] + "fastmcp_tool_call_duration_seconds", "Duration of MCP tool calls by tool.", ["tool"], registry=registry ), ) @@ -281,6 +304,12 @@ class FastMcpPrometheusInstrument(PrometheusInstrument[FastMcpConfig]): _default_metrics: typing.Callable[..., typing.Any] | None = dataclasses.field( default=None, init=False, repr=False, compare=False ) + _inprogress_gauges: dict[tuple[str, tuple[str, ...]], "prometheus_client.Gauge"] = dataclasses.field( + default_factory=dict, init=False, repr=False, compare=False + ) + _inprogress_gauge_lock: LockType = dataclasses.field( + default_factory=threading.Lock, init=False, repr=False, compare=False + ) @staticmethod def dependencies_installed() -> bool: @@ -289,9 +318,12 @@ def dependencies_installed() -> bool: and import_checker.is_prometheus_fastapi_instrumentator_installed ) - def _instrument_http_app(self, http_application: "StarletteWithLifespan") -> "StarletteWithLifespan": + def _instrument_http_app( + self, http_application: "StarletteWithLifespan", *, registry: "prometheus_client.CollectorRegistry" + ) -> "StarletteWithLifespan": config: typing.Final = self.bootstrap_config instrumentator_params = dict(config.prometheus_instrumentator_additional_params) + instrumentator_params["registry"] = registry instrumentator_params["excluded_handlers"] = [ *instrumentator_params.get("excluded_handlers", []), *self.infrastructure_excluded_patterns(), @@ -305,19 +337,52 @@ def _instrument_http_app(self, http_application: "StarletteWithLifespan") -> "St if name in _DEFAULT_METRICS_PARAMETERS } self._default_metrics = instrumentator_metrics.default(**default_metrics_params) - Instrumentator(**instrumentator_params).add(self._default_metrics).instrument( - http_application, **config.prometheus_instrument_additional_params - ) + middleware_count = len(http_application.user_middleware) + instrumentator = Instrumentator(**instrumentator_params).add(self._default_metrics) + instrumentator.instrument(http_application, **config.prometheus_instrument_additional_params) + if ( + registry is not prometheus_client.REGISTRY + and instrumentator.should_instrument_requests_inprogress + and len(http_application.user_middleware) > middleware_count + ): + # The instrumentator registers its built-in gauge globally even with a custom registry. + labels = ("method", "handler") if instrumentator.inprogress_labels else () + gauge = self._get_inprogress_gauge(registry, instrumentator.inprogress_name, labels) + http_application.user_middleware[0] = StarletteMiddleware( + _FastMcpPrometheusHttpMiddleware, + inprogress_gauge=gauge, + **http_application.user_middleware[0].kwargs, + ) return http_application + def _get_inprogress_gauge( + self, registry: "prometheus_client.CollectorRegistry", name: str, labels: tuple[str, ...] + ) -> "prometheus_client.Gauge": + key = (name, labels) + with self._inprogress_gauge_lock: + gauge = self._inprogress_gauges.get(key) + if gauge is None: + gauge = prometheus_client.Gauge( + name, + "Number of HTTP requests in progress.", + labels, + registry=registry, + multiprocess_mode="livesum", + ) + self._inprogress_gauges[key] = gauge + return gauge + def bootstrap(self) -> None: config = self.bootstrap_config + registry = config.prometheus_instrumentator_additional_params.get("registry") or prometheus_client.REGISTRY if config.is_prometheus_http_metrics_on( opentelemetry_records_http_metrics=FastMcpOpenTelemetryInstrument.is_recording_metrics(config) ): - self._restore_http_app = _postprocess_http_apps(config.application, self._instrument_http_app) + self._restore_http_app = _postprocess_http_apps( + config.application, functools.partial(self._instrument_http_app, registry=registry) + ) if config.fastmcp_prometheus_tool_metrics_enabled: - config.application.add_middleware(FastMcpPrometheusMiddleware(*_build_tool_call_metrics())) + config.application.add_middleware(FastMcpPrometheusMiddleware(*_build_tool_call_metrics(registry))) @config.application.custom_route( config.prometheus_metrics_path, @@ -327,7 +392,7 @@ def bootstrap(self) -> None: ) async def metrics_handler(_: "Request") -> "Response": return Response( - prometheus_client.generate_latest(prometheus_client.REGISTRY), + prometheus_client.generate_latest(registry), headers={"content-type": prometheus_client.CONTENT_TYPE_LATEST}, ) diff --git a/tests/test_fastmcp_bootstrap.py b/tests/test_fastmcp_bootstrap.py index b7726be..3134143 100644 --- a/tests/test_fastmcp_bootstrap.py +++ b/tests/test_fastmcp_bootstrap.py @@ -21,6 +21,7 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import SpanKind, get_tracer_provider +from prometheus_client.parser import text_string_to_metric_families from starlette import status from starlette.applications import Starlette from starlette.requests import Request @@ -672,6 +673,99 @@ async def items(_request: Request) -> Response: assert _http_requests_total(path) == expected_requests +def _scraped_samples(text: str) -> list[tuple[str, dict[str, str], float]]: + return [ + (sample.name, sample.labels, sample.value) + for family in text_string_to_metric_families(text) + for sample in family.samples + ] + + +def test_fastmcp_prometheus_scrapes_http_and_tool_metrics_from_custom_registry() -> None: + registry = prometheus_client.CollectorRegistry() + prometheus_client.Counter("custom_requests", "Custom requests.", registry=registry).inc() + config = _make_test_config(prometheus_instrumentator_additional_params={"registry": registry}) + + @config.application.custom_route("/items", methods=["GET"]) + async def items(_request: Request) -> Response: + return Response("ok") + + @config.application.tool + def echo() -> str: + return "ok" + + bootstrapper = FastMcpBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + try: + first = application.http_app() + second = application.http_app(path="/other", stateless_http=True, json_response=True) + with TestClient(first) as client: + assert client.get("/items").is_success + with TestClient(second) as client: + assert client.get("/items").is_success + result = client.post( + "/other", + headers={"Accept": "application/json, text/event-stream"}, + json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "echo"}}, + ) + assert result.is_success + assert result.json()["result"]["content"] == [{"type": "text", "text": "ok"}] + scrape = client.get("/metrics") + finally: + bootstrapper.teardown() + + assert scrape.is_success + samples = _scraped_samples(scrape.text) + assert ("http_requests_total", {"handler": "/items", "method": "GET", "status": "2xx"}, 2) in samples + assert ("fastmcp_tool_calls_total", {"status": "success", "tool": "echo"}, 1) in samples + assert ("fastmcp_tool_call_duration_seconds_count", {"tool": "echo"}, 1) in samples + assert ("custom_requests_total", {}, 1) in samples + + +@pytest.mark.parametrize("change_gauge_options", [False, True]) +def test_fastmcp_prometheus_scrapes_inprogress_gauge_from_custom_registry(change_gauge_options: bool) -> None: + registry = prometheus_client.CollectorRegistry() + gauge_name = f"test_inprogress_{uuid.uuid4().hex}" + second_gauge_name = f"{gauge_name}_second" + labels = {"method": "GET", "handler": "/items"} + config = _make_test_config( + prometheus_instrumentator_additional_params={ + "registry": registry, + "should_instrument_requests_inprogress": True, + "inprogress_name": gauge_name, + "inprogress_labels": True, + }, + ) + + @config.application.custom_route("/items", methods=["GET"]) + async def items(_request: Request) -> Response: + assert 1 in (registry.get_sample_value(gauge_name, labels), registry.get_sample_value(second_gauge_name)) + return Response("ok") + + bootstrapper = FastMcpBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + try: + first = application.http_app() + if change_gauge_options: + config.prometheus_instrumentator_additional_params.update( + inprogress_name=second_gauge_name, inprogress_labels=False + ) + second = application.http_app(path="/other") + with TestClient(first) as first_client, TestClient(second) as second_client: + assert first_client.get("/items").is_success + assert second_client.get("/items").is_success + scrape = second_client.get("/metrics") + finally: + bootstrapper.teardown() + + assert scrape.is_success + samples = _scraped_samples(scrape.text) + assert (gauge_name, labels, 0) in samples + if change_gauge_options: + assert (second_gauge_name, {}, 0) in samples + assert prometheus_client.REGISTRY.get_sample_value(gauge_name, labels) is None + + def test_fastmcp_prometheus_teardown_restores_http_app() -> None: bootstrapper = FastMcpBootstrapper(bootstrap_config=_make_test_config()) application = bootstrapper.bootstrap()