"""Observability primitives (ADR 0012).
This module provides:
* Structured logging initialization (stderr only)
* Correlation ID generation per request
* Minimal in-process metrics counters
* Timing utilities and a no-op trace span abstraction
Design goals:
- Zero external dependencies.
- Safe to import early (lazy initialization where possible).
- Does not write to stdout (stdin/stdout reserved for MCP protocol).
"""
from __future__ import annotations
import json
import logging
import os
import sys
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass
from threading import RLock
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING: # pragma: no cover - type checking only
from collections.abc import Generator
LOG_LEVEL_ENV = "STAC_MCP_LOG_LEVEL"
LOG_FORMAT_ENV = "STAC_MCP_LOG_FORMAT" # "text" | "json"
ENABLE_METRICS_ENV = "STAC_MCP_ENABLE_METRICS"
ENABLE_TRACE_ENV = "STAC_MCP_ENABLE_TRACE"
LATENCY_BUCKETS_ENV = (
"STAC_MCP_LATENCY_BUCKETS_MS" # comma-separated e.g. "5,10,25,50,100,250,500,1000"
)
_logger_state = {"initialized": False}
# Backward compatibility shim: tests (and possibly external code) reference
# observability._logger_initialized. Maintain it as an alias to internal state.
_logger_initialized = False # historical public alias retained (N816 accepted)
_init_lock = RLock()
def _get_bool(env: str, default: bool) -> bool:
val = os.getenv(env)
if val is None:
return default
return val.lower() in {"1", "true", "yes", "on"}
def init_logging() -> None:
"""Configure the library logger (re-initializable for tests).
Tests historically flipped `_logger_initialized = False` to force a fresh
configuration within a stderr/stdout capture context. We preserve that
behavior by allowing re-init when the shim flag is False even if the
internal state dict says initialized.
"""
# Allow re-init when external test code flips alias to False. Avoid using
# global assignment; we rely on shared mutable state and alias pointer.
if _logger_state["initialized"] and _logger_initialized: # pragma: no cover
return
with _init_lock:
if _logger_state["initialized"] and _logger_initialized: # pragma: no cover
return
level_name = os.getenv(LOG_LEVEL_ENV, "WARNING").upper()
level = getattr(logging, level_name, logging.WARNING)
log_format = os.getenv(LOG_FORMAT_ENV, "text").lower()
handler = logging.StreamHandler(stream=sys.stderr)
if log_format == "json":
handler.setFormatter(JSONLogFormatter())
else:
handler.setFormatter(
logging.Formatter("%(levelname)s %(name)s: %(message)s"),
)
logger = logging.getLogger("stac_mcp")
logger.setLevel(level)
logger.handlers = [handler]
logger.propagate = False
_logger_state["initialized"] = True
# Keep alias in sync (tests may introspect this value)
globals()["_logger_initialized"] = True
class JSONLogFormatter(logging.Formatter):
"""Serialize log records as single-line JSON objects (structured logging)."""
def format(self, record: logging.LogRecord) -> str: # type: ignore[override]
base = {
"timestamp": time.strftime(
"%Y-%m-%dT%H:%M:%SZ",
time.gmtime(record.created),
),
"level": record.levelname,
"logger": record.name,
"message": record.getMessage(),
}
# Attach extra known attributes if present
for attr in [
"event",
"tool_name",
"duration_ms",
"error_type",
"correlation_id",
"cache_hit",
"catalog_url",
]:
if hasattr(record, attr):
base[attr] = getattr(record, attr)
return json.dumps(base, separators=(",", ":"))
# ---------------------- Metrics Registry ---------------------- #
class MetricsRegistry:
"""In-process metrics counters + latency histograms (thread-safe)."""
def __init__(self) -> None:
self._lock = RLock()
self._counters: dict[str, int] = {}
# Histogram buckets: parse env or default set (ms)
self._latency_buckets = self._parse_buckets()
# Map metric name -> list[counts per bucket + overflow]
self._histograms: dict[str, list[int]] = {}
def _parse_buckets(self) -> list[float]:
raw = os.getenv(LATENCY_BUCKETS_ENV)
if raw:
try:
buckets = sorted(
{float(x.strip()) for x in raw.split(",") if x.strip()},
)
return [b for b in buckets if b > 0]
except (ValueError, TypeError) as exc: # pragma: no cover - fallback path
logging.getLogger("stac_mcp").debug(
"Invalid STAC_MCP_LATENCY_BUCKETS_MS value: %s",
exc,
)
# Sensible default spanning sub-ms to multi-second
return [1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2000, 5000]
def inc(self, name: str, amount: int = 1) -> None:
if not _get_bool(ENABLE_METRICS_ENV, True):
return
with self._lock:
self._counters[name] = self._counters.get(name, 0) + amount
def observe_latency(self, name: str, value_ms: float) -> None:
if not _get_bool(ENABLE_METRICS_ENV, True): # pragma: no cover - simple guard
return
with self._lock:
hist = self._histograms.get(name)
if hist is None:
hist = [0] * (len(self._latency_buckets) + 1) # last bucket = overflow
self._histograms[name] = hist
# Find first bucket >= value, else overflow
placed = False
for idx, upper in enumerate(self._latency_buckets):
if value_ms <= upper:
hist[idx] += 1
placed = True
break
if not placed: # overflow
hist[-1] += 1
def snapshot(self) -> dict[str, int]:
with self._lock:
return dict(self._counters)
def latency_snapshot(self) -> dict[str, dict[str, int]]:
with self._lock:
snap: dict[str, dict[str, int]] = {}
for name, counts in self._histograms.items():
bucket_labels = [f"<= {int(b)}ms" for b in self._latency_buckets] + [
"> overflow",
]
snap[name] = {bucket_labels[i]: counts[i] for i in range(len(counts))}
return snap
metrics = MetricsRegistry()
def _metric_name(*parts: str) -> str:
return ".".join(parts)
# ---------------------- Tracing (no-op) ---------------------- #
@contextmanager
def trace_span(name: str, **_attrs: Any) -> Generator[None, None, None]:
"""No-op span context manager placeholder.
If tracing is enabled (`STAC_MCP_ENABLE_TRACE`), we could log span start/end.
For now, it's intentionally minimal to avoid overhead.
"""
enabled = _get_bool(ENABLE_TRACE_ENV, False)
t0 = time.perf_counter()
try:
yield
finally: # pragma: no branch - single exit path
if enabled:
duration_ms = (time.perf_counter() - t0) * 1000.0
logging.getLogger("stac_mcp").debug(
"trace_span",
extra={
"event": "trace_span",
"span": name,
"duration_ms": round(duration_ms, 2),
},
)
# ---------------------- Correlation IDs ---------------------- #
def new_correlation_id() -> str:
return str(uuid.uuid4())
@dataclass
class ToolExecutionResult:
"""Container for instrumented tool execution output."""
value: Any
correlation_id: str
duration_ms: float
error_type: str | None = None
def instrument_tool_execution(
tool_name: str,
catalog_url: str | None,
func,
*args,
**kwargs,
) -> ToolExecutionResult:
"""Execute a tool handler with logging, timing, metrics, and correlation id.
Parameters
----------
tool_name: str
Name of the tool being executed.
catalog_url: Optional[str]
Catalog endpoint associated with the execution (may be None).
func: Callable
The handler function to execute.
*args, **kwargs:
Passed to the handler.
"""
init_logging()
correlation_id = new_correlation_id()
logger = logging.getLogger("stac_mcp")
metrics.inc(_metric_name("tool_invocations_total", tool_name))
t0 = time.perf_counter()
error_type: str | None = None
try:
with trace_span(f"tool.{tool_name}"):
result = func(*args, **kwargs)
duration_ms = (time.perf_counter() - t0) * 1000.0
return ToolExecutionResult(
value=result,
correlation_id=correlation_id,
duration_ms=duration_ms,
)
except Exception as exc:
# Classify error type (simple heuristic)
etype = type(exc).__name__
if "timeout" in etype.lower():
error_type = "TimeoutError"
elif "network" in etype.lower() or "connection" in etype.lower():
error_type = "NetworkError"
else:
error_type = "UnknownError"
metrics.inc(_metric_name("tool_errors_total", tool_name, error_type))
duration_ms = (time.perf_counter() - t0) * 1000.0
logger.warning(
"tool_error",
extra={
"event": "tool_error",
"tool_name": tool_name,
"error_type": error_type,
"correlation_id": correlation_id,
"duration_ms": round(duration_ms, 2),
"catalog_url": catalog_url,
},
)
raise
finally:
# Record latency exactly once (success or failure) and log success if applicable
duration_ms = (time.perf_counter() - t0) * 1000.0
metrics.observe_latency(_metric_name("tool_latency_ms", tool_name), duration_ms)
if error_type is None:
logger.info(
"tool_complete",
extra={
"event": "tool_complete",
"tool_name": tool_name,
"duration_ms": round(duration_ms, 2),
"correlation_id": correlation_id,
"catalog_url": catalog_url,
},
)
def metrics_snapshot() -> dict[str, int]:
"""Return a copy of current counter values (for tests)."""
return metrics.snapshot()
def metrics_latency_snapshot() -> dict[str, dict[str, int]]:
"""Return current latency histogram snapshots."""
return metrics.latency_snapshot()
