diff --git a/README.md b/README.md index bf299b1..e21737e 100644 --- a/README.md +++ b/README.md @@ -313,6 +313,29 @@ Example: server = ServerConfig(name="my-server", transport_type="http", max_sessions=100, session_timeout_ms=1_800_000) ``` +### Health resources + +Decorated health checks are exposed at both `health://status` and `health://checks`. +The status resource keeps its original `{check_name: status}` JSON map for existing +clients. The checks resource returns an aggregate payload with `status`, `uptimeSeconds`, +per-check `name`, `status`, `latencyMs`, and `message`, plus `count`, `transforms`, +`timestamp`, and `sdkVersion`. Checks default to a five-second timeout and run independently; +set a per-check timeout with `@health_check("database", timeout=2.0)`. A false check result +marks the aggregate `degraded`; a timeout or exception marks it `unhealthy` (exceptions are +reported with a generic message on `health://checks`; the legacy status resource retains its +existing error string). A timed-out check has status `timeout`. + +Async checks are timed out on the event loop. Sync checks run in a dedicated four-worker pool, +with at most four submitted sync checks at a time; a timed-out sync function keeps its worker +slot until it returns, but does not hold up the resource response or `asyncio.run()` shutdown. +Executor threads are non-daemon, so a sync check that never returns can still delay process exit. +Sync checks therefore run off the main thread and should not depend on thread-bound resources +such as a database connection created on the main thread. Transform entries currently report +their name, type, and order; `details` is reserved for telemetry that is not yet collected. +The SDK version is resolved once when the package loads. + +`ServerHealthPayload` and `TransformTelemetry` are exported from `nitrostack` for typing. + `ServerConfig.protocol_era` is used only when `MCP_STATELESS` and `NITRO_MCP_PROTOCOL_VERSION` are both unset. Unknown tokens become `auto`, same as an unknown env value. --- diff --git a/nitrostack/__init__.py b/nitrostack/__init__.py index 4274c09..9322d9c 100644 --- a/nitrostack/__init__.py +++ b/nitrostack/__init__.py @@ -1,5 +1,26 @@ # NitroStack Python SDK Public API Exports +try: + from importlib.metadata import PackageNotFoundError, version as _package_version + + __version__ = _package_version("nitrostack") +except PackageNotFoundError: + # Keep the SDK version available when running directly from a source checkout. + try: + from pathlib import Path + + try: + import tomllib + except ModuleNotFoundError: # Python 3.10 + import tomli as tomllib + + _project_metadata = tomllib.loads( + (Path(__file__).resolve().parent.parent / "pyproject.toml").read_text(encoding="utf-8") + ) + __version__ = str(_project_metadata["project"]["version"]) + except (ImportError, KeyError, OSError, ValueError, TypeError): + __version__ = "unknown" + from nitrostack.core.decorators import ( tool, resource, @@ -69,6 +90,8 @@ cache, rate_limit, health_check, + TransformTelemetry, + ServerHealthPayload, ) from nitrostack.events.event_emitter import ( on_event, @@ -144,6 +167,7 @@ __all__ = [ + "__version__", "tool", "resource", "prompt", @@ -178,6 +202,8 @@ "cache", "rate_limit", "health_check", + "TransformTelemetry", + "ServerHealthPayload", "on_event", "EventEmitter", "ApiKeyModule", diff --git a/nitrostack/core/additional_decorators.py b/nitrostack/core/additional_decorators.py index 81620eb..b0d91e6 100644 --- a/nitrostack/core/additional_decorators.py +++ b/nitrostack/core/additional_decorators.py @@ -1,7 +1,10 @@ import time import inspect +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor from functools import wraps -from typing import Callable, Any, Dict, List, Tuple +from typing import Callable, Any, Dict, List, Tuple, TypedDict, Literal, Optional from nitrostack.core.context import ExecutionContext def copy_mcp_attributes(src: Any, dst: Any) -> None: @@ -77,28 +80,64 @@ async def wrapper(*args, **kwargs): return decorator +class TransformTelemetry(TypedDict): + name: str + type: str + order: int + details: Dict[str, Any] + + +class HealthCheckResult(TypedDict): + name: str + status: Literal["healthy", "unhealthy", "error", "timeout"] + latencyMs: float + message: Optional[str] + + +class ServerHealthPayload(TypedDict, total=False): + status: Literal["healthy", "degraded", "unhealthy"] + uptimeSeconds: int + checks: List[HealthCheckResult] + count: int + transforms: List[TransformTelemetry] + timestamp: str + sdkVersion: str + + class HealthCheckRegistry: # Maps name -> (unbound_func, class_type) _checks: Dict[str, Tuple[Callable, Any]] = {} # Maps name -> bound_callable - _bound_checks: Dict[str, Callable[[], bool]] = {} + _bound_checks: Dict[str, Tuple[Callable[[], bool], float]] = {} + DEFAULT_TIMEOUT_SECONDS = 5.0 + SYNC_EXECUTOR_WORKERS = 4 + _sync_executor = ThreadPoolExecutor( + max_workers=SYNC_EXECUTOR_WORKERS, + thread_name_prefix="nitrostack-health", + ) + # Bound submitted work as well as running threads. Timed-out functions retain + # their slot until they actually return, preventing an unbounded work queue. + _sync_slots = threading.BoundedSemaphore(SYNC_EXECUTOR_WORKERS) @classmethod def register(cls, name: str, func: Callable, class_type: Any = None) -> None: cls._checks[name] = (func, class_type) @classmethod - def bind_instance(cls, name: str, func: Callable, instance: Any) -> None: + def bind_instance(cls, name: str, func: Callable, instance: Any, timeout: float = 5.0) -> None: # Bind the method to the resolved instance import inspect if inspect.ismethod(func): - cls._bound_checks[name] = func + cls._bound_checks[name] = (func, timeout) else: - cls._bound_checks[name] = lambda: func(instance) + cls._bound_checks[name] = (lambda: func(instance), timeout) @classmethod def get_checks(cls) -> Dict[str, Callable[[], bool]]: - return cls._bound_checks + return { + name: value[0] if isinstance(value, tuple) else value + for name, value in cls._bound_checks.items() + } @classmethod async def run_all(cls) -> Dict[str, str]: @@ -107,25 +146,70 @@ async def run_all(cls) -> Dict[str, str]: Coroutine checks are awaited, sync checks are called directly, and a check that raises reports its own error without affecting the others. """ - results: Dict[str, str] = {} - for name, check in list(cls._bound_checks.items()): + detailed = await cls.run_detailed() + return { + result["name"]: ( + f'error: {result["message"]}' if result["status"] == "error" + else "unhealthy" if result["status"] == "timeout" + else result["status"] + ) + for result in detailed + } + + @classmethod + async def run_detailed(cls) -> List[HealthCheckResult]: + """Run all checks concurrently, bounding each by its own timeout.""" + async def run_one(name: str, check: Callable, timeout: float) -> HealthCheckResult: + started = time.perf_counter() try: - outcome = check() - if inspect.isawaitable(outcome): - outcome = await outcome - results[name] = "healthy" if outcome else "unhealthy" + async def invoke(): + if inspect.iscoroutinefunction(check): + return await check() + while not cls._sync_slots.acquire(blocking=False): + await asyncio.sleep(0.01) + try: + future = cls._sync_executor.submit(check) + except BaseException: + cls._sync_slots.release() + raise + future.add_done_callback(lambda _future: cls._sync_slots.release()) + result = await asyncio.wrap_future(future) + return await result if inspect.isawaitable(result) else result + + outcome = await asyncio.wait_for(invoke(), timeout=timeout) + status = "healthy" if outcome else "unhealthy" + message = None if outcome else "Check returned false" + except asyncio.TimeoutError: + status, message = "timeout", f"Timed out after {timeout:g}s" except Exception as exc: # noqa: BLE001 - one bad check must not hide the rest - results[name] = f"error: {exc}" - return results - -def health_check(name: str): + status, message = "error", str(exc) + return { + "name": name, + "status": status, + "latencyMs": round((time.perf_counter() - started) * 1000, 3), + "message": message, + } + + checks = [ + (name, value if isinstance(value, tuple) else (value, cls.DEFAULT_TIMEOUT_SECONDS)) + for name, value in cls._bound_checks.items() + ] + return list(await asyncio.gather(*( + run_one(name, check, timeout) for name, (check, timeout) in checks + ))) + +def health_check(name: str, timeout: float = HealthCheckRegistry.DEFAULT_TIMEOUT_SECONDS): """ Decorator to mark a service or controller method as a health check. """ + if timeout <= 0: + raise ValueError("health check timeout must be greater than zero") + def decorator(func: Callable): func._mcp_health_check_name = name # We don't have the class type here yet, but we will register it. # It will be bound during discovery in the McpApplicationFactory. HealthCheckRegistry.register(name, func) + func._mcp_health_check_timeout = timeout return func return decorator diff --git a/nitrostack/core/app.py b/nitrostack/core/app.py index 255d504..bf7fbaf 100644 --- a/nitrostack/core/app.py +++ b/nitrostack/core/app.py @@ -6,6 +6,7 @@ import asyncio import inspect import datetime +import time import logging from dataclasses import dataclass, field from pathlib import Path @@ -27,6 +28,7 @@ UnsubscribeRequestParams, ) from pydantic import BaseModel, create_model +from nitrostack import __version__ as _sdk_version from nitrostack.transforms import ( TransformContext, TransformError, @@ -48,7 +50,7 @@ ) from nitrostack.core.mcp_server import NitroStackMcpServer from nitrostack.core.pipeline import run_pipeline -from nitrostack.core.additional_decorators import HealthCheckRegistry +from nitrostack.core.additional_decorators import HealthCheckRegistry, ServerHealthPayload from nitrostack.core.task import TaskManager, TaskStatus from nitrostack.events.event_emitter import EventEmitter from nitrostack.protocol.schema import ( @@ -584,6 +586,8 @@ def _auth_metadata_from_request_ctx(rc: Any) -> Dict[str, Any]: class McpApplication: def __init__(self, app_class: Type): self.app_class = app_class + self._started_at = time.monotonic() + self._sdk_version = _sdk_version if hasattr(app_class, "_mcp_app_module"): self.root_module = app_class._mcp_app_module self.server_config = app_class._mcp_app_server @@ -679,7 +683,10 @@ def _bootstrap(self) -> None: # Bind Health Checks if hasattr(member, "_mcp_health_check_name"): check_name = getattr(member, "_mcp_health_check_name") - HealthCheckRegistry.bind_instance(check_name, member, instance) + HealthCheckRegistry.bind_instance( + check_name, member, instance, + timeout=float(getattr(member, "_mcp_health_check_timeout", 5.0)), + ) # Bind Event Listeners if hasattr(member, "_mcp_event_name"): @@ -829,17 +836,71 @@ def _register_prompt(self, instance: Any, method: Callable, prompt_config: Promp self._prompts[prompt_config.name] = _PromptEntry(config=prompt_config, instance=instance, method=method) def _register_health_resource(self) -> None: - config = ResourceConfig( + status_config = ResourceConfig( uri="health://status", name="Health Status", description="System health status check", + mime_type="application/json", + ) + checks_config = ResourceConfig( + uri="health://checks", + name="Health Checks & Pipeline Telemetry", + description="Current health status and transform pipeline telemetry", + mime_type="application/json", ) async def health_status_resource(context: ExecutionContext) -> str: + # Keep the original string map at health://status for existing clients. results = await HealthCheckRegistry.run_all() return json.dumps(results) - self._resources[config.uri] = _ResourceEntry(config=config, instance=None, method=health_status_resource) + async def health_checks_resource(context: ExecutionContext) -> ServerHealthPayload: + checks = await HealthCheckRegistry.run_detailed() + has_failures = any( + check["status"] in {"error", "timeout"} + for check in checks + ) + has_unhealthy = any( + check["status"] == "unhealthy" + for check in checks + ) + transforms = [] + configured_transforms = getattr(self.app_class, "_mcp_transforms", None) or [] + for order, transform in enumerate(configured_transforms): + transforms.append({ + "name": getattr(transform, "name", transform.__class__.__name__), + "type": transform.__class__.__name__, + "order": order, + "details": {}, + }) + public_checks = [ + { + **check, + "message": "Check failed" if check["status"] == "error" else check["message"], + } + for check in checks + ] + payload: ServerHealthPayload = { + "status": ( + "unhealthy" if has_failures + else "degraded" if has_unhealthy + else "healthy" + ), + "uptimeSeconds": int(time.monotonic() - self._started_at), + "checks": public_checks, + "count": len(public_checks), + "transforms": transforms, + "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z"), + "sdkVersion": self._sdk_version, + } + return payload + + self._resources[status_config.uri] = _ResourceEntry( + config=status_config, instance=None, method=health_status_resource + ) + self._resources[checks_config.uri] = _ResourceEntry( + config=checks_config, instance=None, method=health_checks_resource + ) # ------------------------------------------------------------------ # Protocol handler wiring (owned low-level `mcp.server.lowlevel.Server`) diff --git a/tests/test_health_async_checks.py b/tests/test_health_async_checks.py index 90374d8..b48765f 100644 --- a/tests/test_health_async_checks.py +++ b/tests/test_health_async_checks.py @@ -1,5 +1,6 @@ import asyncio import json +import threading import warnings import pytest @@ -69,4 +70,105 @@ def test_raising_check_reports_error_without_affecting_others(): def test_no_never_awaited_warning(): with warnings.catch_warnings(): warnings.simplefilter("error", RuntimeWarning) - _read_status() \ No newline at end of file + _read_status() + + +def _read_checks(app_module): + async def run(): + harness = await NitroTestingModule.create(app_module) + return await harness.read_resource("health://checks") + + return asyncio.run(run()) + + +def test_health_payload_distinguishes_healthy_degraded_and_failing(): + @injectable() + class Svc: + @health_check("healthy") + def healthy(self): + return True + + @health_check("degraded") + def degraded(self): + return False + + @health_check("failing") + def failing(self): + raise RuntimeError("dependency unavailable") + + @module(name="health-statuses", providers=[Svc]) + class M: + pass + + payload = _read_checks(M) + statuses = {check["name"]: check["status"] for check in payload["checks"]} + assert statuses == {"healthy": "healthy", "degraded": "unhealthy", "failing": "error"} + failing_check = next(check for check in payload["checks"] if check["name"] == "failing") + assert failing_check["message"] == "Check failed" + assert payload["status"] == "unhealthy" + + +def test_false_check_degrades_aggregate_without_failing_it(): + @injectable() + class Svc: + @health_check("partial") + def partial(self): + return False + + @module(name="health-degraded", providers=[Svc]) + class M: + pass + + payload = _read_checks(M) + assert payload["status"] == "degraded" + + +def test_timed_out_sync_check_does_not_block_asyncio_run_or_other_checks(): + release = threading.Event() + + @injectable() + class Svc: + @health_check("slow_sync", timeout=0.05) + def slow_sync(self): + release.wait(timeout=5) + return True + + @health_check("fast") + async def fast(self): + await asyncio.sleep(0.005) + return True + + @module(name="health-timeouts", providers=[Svc]) + class M: + pass + + result = {} + read_complete = threading.Event() + + def read_resource(): + try: + result["payload"] = _read_checks(M) + finally: + read_complete.set() + + reader = threading.Thread(target=read_resource) + reader.start() + completed_before_release = read_complete.wait(timeout=2) + release.set() + reader.join(timeout=2) + + assert completed_before_release, "asyncio.run waited for the timed-out sync check" + assert not reader.is_alive() + payload = result["payload"] + statuses = {check["name"]: check for check in payload["checks"]} + assert statuses["slow_sync"]["status"] == "timeout" + assert statuses["slow_sync"]["message"] == "Timed out after 0.05s" + assert statuses["fast"]["status"] == "healthy" + assert payload["status"] == "unhealthy" + + +def test_health_status_retains_legacy_string_map(): + status = _read_status() + assert status["async_ok"] == "healthy" + assert status["sync_ok"] == "healthy" + assert status["async_boom"] == "error: db down"