mcp-telemetry 0.2.0__py3-none-any.whl

mcp_telemetry/__init__.py

@@ -0,0 +1,41 @@
+"""mcp-telemetry — zero-config observability for MCP servers.
+
+Core is stdlib-only. Optional exporters:
+  - OTLP/HTTP (needs httpx)  -> pip install mcp-telemetry[otlp]
+  - built-in JSONL + live TUI terminal panel (no deps)
+
+Extended surface — propagation, sampling, metrics, OTel provider, replay:
+  from mcp_telemetry.propagator import generate_traceparent, parse_traceparent
+  from mcp_telemetry.sampler import parent_based, ratio, rate_limited
+  from mcp_telemetry.metrics import Registry, metrics_from_store
+  from mcp_telemetry.otel_provider import OtelProvider
+  from mcp_telemetry.fastmcp import patch_fastmcp, make_server
+  mcp-trace --replay store.jsonl --console
+"""
+from mcp_telemetry.api import (
+    auto,
+    instrument,
+    session,
+    span,
+    trace_store,
+)
+from mcp_telemetry import monkey
+from mcp_telemetry.monkey import wrap_tool_call, patch_mcp_sdk
+from mcp_telemetry.store import TraceStore, default_store, set_default_store
+
+__all__ = [
+    "auto",
+    "instrument",
+    "session",
+    "span",
+    "trace_store",
+    "monkey",
+    "wrap_tool_call",
+    "patch_mcp_sdk",
+    "TraceStore",
+    "default_store",
+    "set_default_store",
+    "__version__",
+]
+
+__version__ = "0.2.0"

mcp_telemetry/api.py

@@ -0,0 +1,98 @@
+"""Public API for mcp-telemetry.
+
+  import mcp_telemetry as mt
+  mt.auto()                      # patch official mcp SDK + fanout to JSONL
+  mt.instrument(exporter=..., scrubbed=True)
+
+  with mt.span("chat"):          # manual span around an LLM call
+      ...
+"""
+from __future__ import annotations
+
+import os
+import sys
+from contextlib import contextmanager
+from typing import Iterator, Optional
+
+from mcp_telemetry import context as _ctx
+from mcp_telemetry import exporters
+from mcp_telemetry import model
+from mcp_telemetry import redact
+from mcp_telemetry.monkey import is_patched, patch_mcp_sdk, wrap_tool_call
+from mcp_telemetry.store import TraceStore, default_store, set_default_store
+
+_ENABLED = False
+
+
+def instrument(
+    store: Optional[TraceStore] = None,
+    trace_exporter: Optional[exporters.Exporter] = None,
+    enabled: bool = True,
+) -> None:
+    """Configure the active store + exporter fanout."""
+    global _ENABLED
+    eff = store or default_store()
+    if store is not None:
+        set_default_store(store)
+    _ENABLED = enabled
+    eff.on_record = exporters.DEFAULT.export if enabled else None
+    if trace_exporter is not None:
+        exporters.DEFAULT._targets = [trace_exporter]
+
+
+def auto(store: Optional[TraceStore] = None, enabled: bool = True) -> bool:
+    """All-in-one: enforce defaults + patch the official mcp SDK.
+
+    If the mcp SDK is missing, patch is a no-op and auto() still returns True
+    so callers can instrument() without knowing the runtime shape.
+    """
+    instrument(store=store, enabled=enabled)
+    if enabled and "MCP_TELEMETRY_DISABLE" not in os.environ:
+        patch_mcp_sdk(store)
+    if enabled and sys.stderr.isatty():
+        exporters.DEFAULT.add(exporters.TextExporter())
+    return patch_mcp_sdk(store) or True
+
+
+def session(trace_id: Optional[str] = None):
+    """Context manager: start a trace, attach all nested spans to it."""
+    store = default_store()
+    trace = store.start(trace_id)
+    prev = _ctx.get_current()
+    _ctx.set_current(trace)
+
+    @contextmanager
+    def _scope() -> Iterator["TraceStore"]:
+        try:
+            yield store
+        finally:
+            store.record(trace)
+            _ctx.set_current(prev)
+
+    return _scope()
+
+
+@contextmanager
+def span(name: str, kind: str = model.SPAN_AGENT) -> Iterator[model.Span]:
+    """Manual span. Attaches to the ambient session trace if one is open,
+    otherwise starts a standalone trace."""
+    store = default_store()
+    current = _ctx.get_current()
+    if current is not None:
+        trace = current
+    else:
+        trace = store.start()
+    sp = trace._make_span(name, None, kind=kind)
+    try:
+        yield sp
+    except Exception as exc:
+        sp.finish("ERROR", f"{type(exc).__name__}: {exc}")
+        store.record(trace, finished=current is None)
+        raise
+    else:
+        sp.finish("OK")
+        store.record(trace, finished=current is None)
+
+
+def trace_store() -> TraceStore:
+    return default_store()

mcp_telemetry/cli.py

@@ -0,0 +1,90 @@
+"""CLI: mcp-trace — render, tail, and clear the local telemetry feed.
+
+  mcp-trace            # last 25 spans in ANSI table
+  mcp-trace --tail     # follow the JSONL feed
+  mcp-trace --json     # dump raw JSONL to stdout (pipe friendly)
+  mcp-trace --clear    # truncate the local JSONL log
+  mcp-trace --replay <file> [--console] [--max N]   # replay recorded traces through the exporter stack
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import time
+
+
+def _read(path: str) -> list[dict]:
+    rows: list[dict] = []
+    try:
+        with open(path, encoding="utf-8") as fh:
+            for line in fh:
+                line = line.strip()
+                if line:
+                    rows.append(json.loads(line))
+    except FileNotFoundError:
+        pass
+    return rows
+
+
+def _render(rows: list[dict]) -> None:
+    for tr in rows[-25:]:
+        for s in tr["spans"]:
+            hit = "OK" if s["status"] == "OK" else "ERR"
+            color = 32 if hit == "OK" else 31
+            print(
+                f"\x1b[36m{tr['id']}\x1b[0m "
+                f"\x1b[1m{s['tool']}\x1b[0m "
+                f"\x1b[{color}m{hit}\x1b[0m "
+                f"\x1b[33m{s.get('latency_ms') or 0:.1f}ms\x1b[0m "
+                f"in#{s.get('input_hash') or '-'} "
+                + (f"\x1b[31m{s['error']}\x1b[0m" if s.get("error") else "")
+            )
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(prog="mcp-trace", description="Render the mcp-telemetry local feed")
+    parser.add_argument("--tail", action="store_true", help="follow the JSONL log")
+    parser.add_argument("--json", action="store_true", help="dump raw records as NDJSON")
+    parser.add_argument("--clear", action="store_true", help="clear the local log")
+    parser.add_argument("--replay", metavar="FILE", help="replay a recorded JSONL store through the exporter stack")
+    parser.add_argument("--console", action="store_true", help="with --replay, print OTLP-shaped traces")
+    parser.add_argument("--max", type=int, metavar="N", help="with --replay, cap the number of traces")
+    parser.add_argument("--path", default="mcp-telemetry.jsonl", help="log file path")
+    args = parser.parse_args()
+
+    if args.replay is not None:
+        from mcp_telemetry.replay import main as replay_main
+
+        sub = [args.replay]
+        if args.console:
+            sub.append("--console")
+        if args.max:
+            sub += ["--max", str(args.max)]
+        return replay_main(sub)
+
+    if args.clear:
+        open(args.path, "w", encoding="utf-8").close()  # noqa: SIM115
+        return 0
+
+    if args.tail:
+        seen = 0
+        while True:
+            rows = _read(args.path)
+            for tr in rows[seen:]:
+                for s in tr["spans"]:
+                    print(json.dumps({"trace": tr["id"], **s}))
+                seen += 1
+            time.sleep(0.5)
+        return 0
+
+    rows = _read(args.path)
+    if args.json:
+        for tr in rows:
+            print(json.dumps(tr))
+        return 0
+    _render(rows)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mcp_telemetry/context.py

@@ -0,0 +1,17 @@
+"""Ambient trace cursor shared across api.py and monkey.py without cycles."""
+from __future__ import annotations
+
+from typing import Optional
+
+from mcp_telemetry.model import Trace
+
+_current: Optional[Trace] = None
+
+
+def set_current(trace: Optional[Trace]) -> None:
+    global _current
+    _current = trace
+
+
+def get_current() -> Optional[Trace]:
+    return _current

mcp_telemetry/exporters.py

@@ -0,0 +1,117 @@
+"""Span payload builders + exporters.
+
+Exporters implement `dump(trace) -> str` so they compose cleanly:
+  - JsonlExporter  -> ndjson for pipes / local debug
+  - OtlpExporter   -> OTLP/HTTP JSON payload (httpx optional extra)
+  - TextExporter   -> ANSI table for the live TUI panel (stdlib only)
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+import threading
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from mcp_telemetry.model import Trace
+
+
+class Exporter(Protocol):
+    def export(self, trace: Trace) -> None: ...
+
+
+@dataclass
+class JsonlExporter:
+    path: str = "mcp-telemetry.jsonl"
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def export(self, trace: Trace) -> None:
+        with self._lock:
+            with open(self.path, "a", encoding="utf-8") as fh:
+                fh.write(json.dumps(trace.to_dict()) + "\n")
+
+
+@dataclass
+class TextExporter:
+    """Minimal ANSI 256-color renderer. No external deps."""
+
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def export(self, trace: Trace) -> None:
+        rows = []
+        for s in trace.spans:
+            status = "\x1b[32mOK\x1b[0m" if s.status == "OK" else f"\x1b[31m{s.status}\x1b[0m"
+            lines = [
+                f"  \x1b[36m{trace.id}\x1b[0m {s.tool or s.name}  {status}  "
+                f"\x1b[33m{s.latency_ms or 0:.1f}ms\x1b[0m"
+            ]
+            if s.input_hash:
+                lines.append(f"    input#{s.input_hash}  {s.server or ''}")
+            if s.error:
+                lines.append(f"    \x1b[31m{s.error}\x1b[0m")
+            rows.extend(lines)
+        payload = "\n".join(rows)
+        with self._lock:
+            print(payload, file=sys.stderr)
+
+
+@dataclass
+class OtlpExporter:
+    """OTLP/HTTP JSON (v1 logs/bridge). Requires httpx (extra: [otlp]).
+
+    Sends one span record per tool call with gen_ai keys in attributes.
+    """
+
+    endpoint: str = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318/v1/logs")
+    service_name: str = os.environ.get("OTEL_SERVICE_NAME", "mcp-server")
+    headers: dict[str, str] = field(default_factory=lambda: {"Content-Type": "application/json"})
+
+    def export(self, trace: Trace) -> None:
+        try:
+            import httpx  # type: ignore
+        except ImportError as exc:  # pragma: no cover
+            raise RuntimeError("OTLP export requires httpx: pip install 'mcp-telemetry[otlp]'") from exc
+
+        records = []
+        for s in trace.spans:
+            attrs = {
+                "gen_ai.tool.name": s.tool or s.name,
+                "gen_ai.server.name": s.server or "",
+                "gen_ai.input.bytes": s.input_hash or "",
+                "gen_ai.usage.input_tokens": s.tokens_in or 0,
+                "gen_ai.usage.output_tokens": s.tokens_out or 0,
+                "gen_ai.operation.span.id": trace.id,
+                "gen_ai.operation.name": s.name,
+            }
+            if s.error:
+                attrs["gen_ai.error.type"] = s.error
+            records.append(
+                {
+                    "scope": {"name": "mcp-telemetry", "version": "0.1.0"},
+                    "resource": {"service.name": self.service_name},
+                    "severity_text": "ERROR" if s.error else "INFO",
+                    "body": {"string_value": f"{s.tool} {s.status}"},
+                    "attributes": attrs,
+                    "timeUnixNano": str(int((s.end or s.start) * 1_000_000_000)),
+                }
+            )
+        httpx.post(self.endpoint, json={"resourceLogs": [{"scopeLogs": [{"logRecords": records}]}]}, headers=self.headers, timeout=5.0)
+
+
+class Fanout:
+    def __init__(self) -> None:
+        self._targets: list[Exporter] = [JsonlExporter()]
+
+    def add(self, exporter: Exporter) -> None:
+        self._targets.append(exporter)
+
+    def export(self, trace: Trace) -> None:
+        for t in self._targets:
+            try:
+                t.export(trace)
+            except Exception:
+                pass
+
+
+DEFAULT = Fanout()

mcp_telemetry/fastmcp.py

@@ -0,0 +1,54 @@
+"""FastMCP integration — one-call instrumentation.
+
+`patch_fastmcp(server)` wraps every tool handler the FastMCP server registers
+so each invocation emits a GenAI-style trace span, deduped per tool, feeding
+the active TraceStore/provider. No SDK import required to USE this module; the
+SDK is fetched lazily so `pip install mcp-telemetry` alone is harmless.
+
+If you do not run FastMCP, call `patch_mcp_sdk()` from the core to auto-wrap
+the stdlib/sdk tool handlers instead — the same spans drop out underneath.
+"""
+from __future__ import annotations
+
+from typing import Any, Callable, Optional
+
+from mcp_telemetry.api import span
+from mcp_telemetry.monkey import wrap_tool_call
+
+
+def patch_fastmcp(server: Any, *, store: Any = None) -> int:
+    """Wrap all handlers currently registered on a FastMCP server object. The
+    instance may be post-`@server.tool(...)` decorated — we walk `server._tool_manager`."""
+    try:
+        import fastmcp  # noqa: F401
+    except ImportError:
+        raise RuntimeError("fastmcp not installed — import it or pip install fastmcp")
+
+    tool_manager = getattr(server, "_tool_manager", None)
+    if tool_manager is None:
+        return 0
+    tools = getattr(tool_manager, "_tools", None) or getattr(tool_manager, "tools", {})
+    wrapped = 0
+    for tool in tools.values() if isinstance(tools, dict) else tools:
+        fn = getattr(tool, "fn", None) or getattr(tool, "func", None)
+        if fn is None:
+            continue
+        name = getattr(tool, "name", None) or getattr(fn, "__name__", "tool")
+        tool.fn = wrap_tool_call(fn, store=store, tool_name=name)  # type: ignore[attr-defined]
+        wrapped += 1
+    return wrapped
+
+
+def make_server(name: str, store: Any) -> Any:
+    """Slim FastMCP-compatible shim for a server name: emits a trace span per
+    tool invocation and records it. Use when you want instrumentation without
+    the SDK present — identical span shape, so replays report the same way."""
+    # class bodies do not close over enclosing function locals — build via type()
+    def tool(fn: Callable[..., Any], *, tool_name: Optional[str] = None) -> Callable[..., Any]:
+        return wrap_tool_call(fn, name=tool_name or getattr(fn, "__name__", "tool"), store=store)
+
+    shim = type("FastMcpShim", (), {"name": name, "tool": staticmethod(tool)})()
+    return shim
+
+
+__all__ = ["patch_fastmcp", "make_server"]

mcp_telemetry/metrics.py

@@ -0,0 +1,157 @@
+"""In-process metrics registry — counters + latency histograms, zero-dep.
+
+Aggregates tool-call counts and latency distributions in memory and exposes:
+  - snapshot()            plain dict for JSON / admin dashboards
+  - openmetrics()         Prometheus/OpenTelemetry text exposition
+  - metrics_from_store()  auto-wire a registry to a TraceStore via on_record
+Render negative-space: no units, no timestamps — just numbers, honestly kept.
+"""
+from __future__ import annotations
+
+import threading
+from typing import Any, Optional
+
+_BUCKETS_MS = (1, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000)
+
+# fine-grain buckets for UI-place histograms; coarser for exporter cost
+_BUCKETS_COARSE = (10, 50, 100, 500, 1000, 5000)
+
+
+class Counter:
+    def __init__(self, name: str, labels: Optional[dict[str, str]] = None) -> None:
+        self.name = name
+        self.labels = labels or {}
+        self.value = 0
+
+    def inc(self, amount: int = 1) -> None:
+        self.value += amount
+
+
+class Histogram:
+    def __init__(self, name: str, labels: Optional[dict[str, str]] = None, buckets: tuple[float, ...] = _BUCKETS_MS) -> None:
+        self.name = name
+        self.labels = labels or {}
+        self.buckets = buckets
+        self.counts = [0] * (len(buckets) + 1)
+        self.sum = 0.0
+        self.count = 0
+        self.min = float("inf")
+        self.max = 0.0
+
+    def observe(self, value_ms: float) -> None:
+        self.sum += value_ms
+        self.count += 1
+        if value_ms < self.min:
+            self.min = value_ms
+        if value_ms > self.max:
+            self.max = value_ms
+        for i, b in enumerate(self.buckets):
+            if value_ms <= b:
+                self.counts[i] += 1
+                break
+        else:
+            self.counts[-1] += 1
+
+    def snapshot(self) -> dict[str, Any]:
+        cumulative = 0
+        cdf: dict[str, int] = {}
+        for i, b in enumerate(self.buckets):
+            cumulative += self.counts[i]
+            cdf[f"le_{b:g}"] = cumulative
+        cdf["le_+Inf"] = self.count
+        p50 = self.quantile(0.5)
+        p95 = self.quantile(0.95)
+        p99 = self.quantile(0.99)
+        return {
+            "count": self.count,
+            "sum_ms": round(self.sum, 3),
+            "min_ms": 0.0 if self.min == float("inf") else round(self.min, 3),
+            "p50_ms": round(p50, 3),
+            "p95_ms": round(p95, 3),
+            "p99_ms": round(p99, 3),
+            "max_ms": round(self.max, 3),
+            "cdf": cdf,
+        }
+
+    def quantile(self, q: float) -> float:
+        if self.count == 0:
+            return 0.0
+        target = self.count * q
+        acc = 0
+        bounds = (0.0,) + self.buckets + (float("inf"),)
+        for i, hi in enumerate(bounds):
+            acc += self.counts[i]
+            if acc >= target:
+                lo = bounds[i]
+                if hi == float("inf"):
+                    return lo
+                return lo + (hi - lo) * 0.5
+        return 0.0
+
+
+class Registry:
+    """Thread-safe metrics registry."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._counters: list[Counter] = []
+        self._histograms: list[Histogram] = []
+
+    def counter(self, name: str, labels: Optional[dict[str, str]] = None) -> Counter:
+        with self._lock:
+            c = Counter(name, labels)
+            self._counters.append(c)
+            return c
+
+    def histogram(self, name: str, labels: Optional[dict[str, str]] = None) -> Histogram:
+        with self._lock:
+            h = Histogram(name, labels)
+            self._histograms.append(h)
+            return h
+
+    def snapshot(self) -> dict[str, Any]:
+        with self._lock:
+            counters = [{"name": c.name, "labels": c.labels, "value": c.value} for c in self._counters]
+            histos = [{"name": h.name, "labels": h.labels, **h.snapshot()} for h in self._histograms]
+        return {"counters": counters, "histograms": histos}
+
+    def openmetrics(self) -> str:
+        """Prometheus text exposition (0.0.1)."""
+        lines = ["# TYPE mcp_metric counter", "# TYPE mcp_latency histogram"]
+        snap = self.snapshot()
+        for c in snap["counters"]:
+            labels = ",".join(f'{k}="{v}"' for k, v in c["labels"].items())
+            suffix = f"{{{labels}}}" if labels else ""
+            lines.append(f'mcp_metric{suffix} {c["value"]}')
+        for h in snap["histograms"]:
+            base = ",".join(f'{k}="{v}"' for k, v in h["labels"].items())
+            for le, val in h["cdf"].items():
+                lbl = f"{base},le=\"{le.removeprefix('le_')}\"" if base else f"le=\"{le.removeprefix('le_')}\""
+                lines.append(f'mcp_latency_bucket{{{lbl}}} {val}')
+            suffix = f"{{{base}}}" if base else ""
+            lines.append(f'mcp_latency_sum{suffix} {h["sum_ms"]}')
+            lines.append(f'mcp_latency_count{suffix} {h["count"]}')
+        return "\n".join(lines) + "\n"
+
+
+def metrics_from_store(store: Any) -> Registry:
+    """Wire a registry to a TraceStore's on_record hook: one counter per tool
+    status, one latency histogram per tool."""
+    reg = Registry()
+    histos: dict[str, Histogram] = {}
+    counters: dict[str, Counter] = {}
+
+    def on_record(trace: Any) -> None:
+        for span in trace.spans:
+            if not span.tool:
+                continue
+            key = span.tool
+            if key not in histos:
+                histos[key] = reg.histogram("mcp_tool_latency", {"tool": key})
+                counters[key] = reg.counter("mcp_calls", {"tool": key, "status": "OK"})
+            histos[key].observe(float(span.latency_ms or 0))
+            counters[key].inc()
+            counters[key].labels["status"] = span.status
+
+    store.on_record = on_record  # type: ignore[attr-defined]
+    return reg

mcp_telemetry/model.py

@@ -0,0 +1,74 @@
+"""Trace/span model aligned with OpenTelemetry GenAI conventions.
+
+Span kinds we emit, matching the OTel GenAI semantic conventions:
+  gen_ai.client.tool_call  -> tool invocation span
+  gen_ai.client.op         -> chat / invoke_agent spans
+
+Every span carries: name in walrus form, stable kinds for the tool layer,
+and an input *fingerprint* (hash) instead of raw payloads by default.
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field, asdict
+from typing import Any, Optional
+
+SPAN_TOOL = "gen_ai.client.tool_call"
+SPAN_CHAT = "gen_ai.client.op"
+SPAN_AGENT = "gen_ai.agent.invoke"
+
+t0 = time.time()
+
+
+def _wall() -> float:
+    return time.time() - t0
+
+
+@dataclass
+class Span:
+    name: str
+    kind: str = SPAN_TOOL
+    tool: Optional[str] = None
+    server: Optional[str] = None
+    start: float = field(default_factory=_wall)
+    end: Optional[float] = None
+    status: str = "OK"
+    error: Optional[str] = None
+    input_hash: Optional[str] = None
+    latency_ms: Optional[float] = None
+    tokens_in: Optional[int] = None
+    tokens_out: Optional[int] = None
+    meta: dict[str, Any] = field(default_factory=dict)
+    parent_span_id: Optional[str] = None
+
+    def finish(self, status: str = "OK", error: Optional[str] = None) -> "Span":
+        self.end = _wall()
+        self.latency_ms = round((self.end - self.start) * 1000, 3)
+        self.status = status
+        self.error = error
+        return self
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class Trace:
+    id: str
+    started: float = field(default_factory=_wall)
+    spans: list[Span] = field(default_factory=list)
+
+    def record(self, span: Span) -> None:
+        self.spans.append(span)
+
+    def _make_span(self, tool: str, server: Optional[str] = None, kind: str = SPAN_TOOL) -> Span:
+        span = Span(name=f"{tool} {kind}", tool=tool, server=server, kind=kind)
+        self.record(span)
+        return span
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"id": self.id, "started": round(self.started, 3), "spans": [s.to_dict() for s in self.spans]}
+
+
+def new_trace(trace_id: Optional[str] = None) -> Trace:
+    return Trace(id=trace_id or f"{int(time.time_ns())}")