Use the configured Prometheus registry in FastMCP

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

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,20 +112,25 @@ 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"]),
             typing.cast("prometheus_client.Histogram", registered_collectors["fastmcp_tool_call_duration_seconds"]),
         )
     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},
             )
 

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()