spanforge 2.0.0__py3-none-any.whl

spanforge/metrics_export.py

@@ -0,0 +1,342 @@
+"""spanforge.metrics_export — Prometheus-compatible metrics export.
+
+This module provides a zero-dependency Prometheus text-format metrics
+endpoint for SpanForge.  It exposes key observability indicators as gauges
+and counters compatible with any Prometheus scraper.
+
+Exported metrics
+----------------
+
+=========================================  =====================================
+Metric name                                Description
+=========================================  =====================================
+``spanforge_spans_total``                  Total spans emitted (counter).
+``spanforge_spans_error_total``            Total error spans (counter).
+``spanforge_export_errors_total``          Total export backend errors (counter).
+``spanforge_events_dropped_total``         Total events dropped (counter).
+``spanforge_token_usage_total``            Total tokens used (counter by type).
+``spanforge_span_duration_ms``             Span duration histogram buckets (gauge).
+``spanforge_drift_alerts_total``           Total drift alerts emitted (counter).
+=========================================  =====================================
+
+Usage
+-----
+Standalone HTTP server::
+
+    from spanforge.metrics_export import serve_metrics
+    serve_metrics(port=9090)   # starts a background thread
+
+Single scrape (e.g. push-gateway integration)::
+
+    from spanforge.metrics_export import PrometheusMetricsExporter, MetricsSummary
+
+    exporter = PrometheusMetricsExporter()
+    text = exporter.export(MetricsSummary(spans_total=1000, error_spans=12, ...))
+    print(text)
+"""
+
+from __future__ import annotations
+
+import http.server
+import logging
+import re
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+__all__ = [
+    "MetricsSummary",
+    "PrometheusMetricsExporter",
+    "serve_metrics",
+]
+
+_log = logging.getLogger("spanforge.metrics_export")
+
+
+# ---------------------------------------------------------------------------
+# MetricsSummary
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class MetricsSummary:
+    """Snapshot of SpanForge observability counters.
+
+    Instances of this class are passed to
+    :meth:`PrometheusMetricsExporter.export` to generate the Prometheus text
+    payload.  All fields default to zero so callers can omit unknown values.
+
+    Args:
+        spans_total:        Cumulative number of spans started.
+        error_spans:        Cumulative number of spans with status ``"error"``.
+        export_errors:      Cumulative export backend errors.
+        events_dropped:     Events silently dropped (queue full / circuit open).
+        prompt_tokens:      Cumulative prompt token count.
+        completion_tokens:  Cumulative completion token count.
+        total_tokens:       Cumulative total token count.
+        total_cost_usd:     Cumulative estimated cost in USD.
+        drift_alerts:       Cumulative drift alert events emitted.
+        active_spans:       Gauge — currently open spans.
+        duration_buckets:   Histogram bucket counts ``{le_ms: count}``.
+        labels:             Optional extra label key/value pairs applied to
+                            every metric (e.g. ``{"service": "my-service"}``).
+        timestamp_ms:       Unix timestamp (milliseconds) of the snapshot.
+    """
+
+    spans_total: int = 0
+    error_spans: int = 0
+    export_errors: int = 0
+    events_dropped: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    total_cost_usd: float = 0.0
+    drift_alerts: int = 0
+    active_spans: int = 0
+    duration_buckets: dict[float, int] = field(default_factory=dict)
+    labels: dict[str, str] = field(default_factory=dict)
+    timestamp_ms: int = field(default_factory=lambda: int(time.time() * 1000))
+
+
+# ---------------------------------------------------------------------------
+# PrometheusMetricsExporter
+# ---------------------------------------------------------------------------
+
+_DEFAULT_DURATION_BUCKETS = (5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 5000.0)
+
+
+class PrometheusMetricsExporter:
+    """Render a :class:`MetricsSummary` as Prometheus text format (0.0.4).
+
+    The output is compatible with Prometheus scraping and the OpenMetrics
+    exposition format.
+
+    Args:
+        namespace: Optional metric name prefix.  Defaults to ``"spanforge"``.
+
+    Example::
+
+        exporter = PrometheusMetricsExporter()
+        summary = MetricsSummary(spans_total=500, error_spans=3)
+        print(exporter.export(summary))
+    """
+
+    def __init__(self, namespace: str = "spanforge") -> None:
+        self._ns = namespace.rstrip("_")
+
+    def export(self, summary: MetricsSummary) -> str:
+        """Return Prometheus text exposition for *summary*.
+
+        Args:
+            summary: Populated :class:`MetricsSummary` snapshot.
+
+        Returns:
+            Multi-line string in Prometheus text format 0.0.4.
+        """
+        ns = self._ns
+        lines: list[str] = []
+        ts = summary.timestamp_ms
+        base_labels = self._format_labels(summary.labels)
+
+        def counter(name: str, help_text: str, value: int | float) -> None:
+            full = f"{ns}_{name}"
+            lines.append(f"# HELP {full} {help_text}")
+            lines.append(f"# TYPE {full} counter")
+            lines.append(f"{full}{base_labels} {value} {ts}")
+
+        def gauge(name: str, help_text: str, value: int | float) -> None:
+            full = f"{ns}_{name}"
+            lines.append(f"# HELP {full} {help_text}")
+            lines.append(f"# TYPE {full} gauge")
+            lines.append(f"{full}{base_labels} {value} {ts}")
+
+        # Span counters
+        counter("spans_total", "Total number of spans emitted.", summary.spans_total)
+        counter("spans_error_total", "Total number of error spans.", summary.error_spans)
+        counter("export_errors_total", "Total export backend errors.", summary.export_errors)
+        counter("events_dropped_total", "Total events dropped.", summary.events_dropped)
+        counter("drift_alerts_total", "Total drift alerts emitted.", summary.drift_alerts)
+
+        # Token usage (with token_type label)
+        tok_name = f"{ns}_token_usage_total"
+        lines.append(f"# HELP {tok_name} Total token usage by token type.")
+        lines.append(f"# TYPE {tok_name} counter")
+        for ttype, count in [
+            ("prompt", summary.prompt_tokens),
+            ("completion", summary.completion_tokens),
+            ("total", summary.total_tokens),
+        ]:
+            label_str = self._format_labels({**summary.labels, "token_type": ttype})
+            lines.append(f"{tok_name}{label_str} {count} {ts}")
+
+        # Cost
+        counter("cost_usd_total", "Total estimated cost in USD.", summary.total_cost_usd)
+
+        # Active spans (gauge)
+        gauge("active_spans", "Currently open (in-flight) spans.", summary.active_spans)
+
+        # Duration histogram
+        if summary.duration_buckets:
+            hist_name = f"{ns}_span_duration_ms"
+            lines.append(f"# HELP {hist_name} Span duration distribution in milliseconds.")
+            lines.append(f"# TYPE {hist_name} histogram")
+            cumulative = 0
+            sorted_buckets = sorted(summary.duration_buckets.items())
+            for le, count in sorted_buckets:
+                cumulative += count
+                le_label = self._format_labels({**summary.labels, "le": str(le)})
+                lines.append(f"{hist_name}_bucket{le_label} {cumulative} {ts}")
+            # +Inf bucket
+            inf_label = self._format_labels({**summary.labels, "le": "+Inf"})
+            lines.append(f"{hist_name}_bucket{inf_label} {cumulative} {ts}")
+        else:
+            # Emit default zero buckets so scrapers don't see missing series.
+            hist_name = f"{ns}_span_duration_ms"
+            lines.append(f"# HELP {hist_name} Span duration distribution in milliseconds.")
+            lines.append(f"# TYPE {hist_name} histogram")
+            for le in _DEFAULT_DURATION_BUCKETS:
+                le_label = self._format_labels({**summary.labels, "le": str(le)})
+                lines.append(f"{hist_name}_bucket{le_label} 0 {ts}")
+            inf_label = self._format_labels({**summary.labels, "le": "+Inf"})
+            lines.append(f"{hist_name}_bucket{inf_label} 0 {ts}")
+
+        lines.append("")  # trailing newline
+        return "\n".join(lines)
+
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _format_labels(labels: dict[str, str]) -> str:
+        if not labels:
+            return ""
+        # Drop any label keys that don't conform to the Prometheus data model.
+        valid_labels = {
+            k: v for k, v in labels.items() if _PROM_LABEL_NAME_RE.match(k)
+        }
+        if not valid_labels:
+            return ""
+        pairs = ",".join(
+            f'{k}="{_escape_label_value(v)}"' for k, v in sorted(valid_labels.items())
+        )
+        return "{" + pairs + "}"
+
+
+# M6: Prometheus label names must match [a-zA-Z_:][a-zA-Z0-9_:]* (Prometheus data model).
+_PROM_LABEL_NAME_RE: re.Pattern[str] = re.compile(r"^[a-zA-Z_:][a-zA-Z0-9_:]*$")
+
+
+def _escape_label_value(value: str) -> str:
+    return value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
+
+
+# ---------------------------------------------------------------------------
+# Live metrics collector — reads from _stream internals
+# ---------------------------------------------------------------------------
+
+
+def _collect_live_summary() -> MetricsSummary:
+    """Build a :class:`MetricsSummary` from live SpanForge stream counters."""
+    summary = MetricsSummary()
+    try:
+        from spanforge._stream import _export_error_count  # noqa: PLC0415
+        summary.export_errors = _export_error_count
+    except Exception:  # NOSONAR
+        pass
+    try:
+        from spanforge._span import _SPAN_STACK  # noqa: PLC0415
+        # _SPAN_STACK is a ContextVar[list]; counting open spans is tricky
+        # without a global registry.  Use 0 as a safe default.
+        _ = _SPAN_STACK
+    except Exception:  # NOSONAR
+        pass
+    return summary
+
+
+# ---------------------------------------------------------------------------
+# HTTP handler + server
+# ---------------------------------------------------------------------------
+
+
+class _MetricsHTTPHandler(http.server.BaseHTTPRequestHandler):
+    """Minimal HTTP handler serving /metrics in Prometheus text format."""
+
+    _exporter: PrometheusMetricsExporter
+    _collector: Any  # callable: () -> MetricsSummary
+
+    def do_GET(self) -> None:  # noqa: N802
+        if self.path != "/metrics":
+            self.send_response(404)
+            self.end_headers()
+            self.wfile.write(b"Not Found\n")
+            return
+
+        try:
+            summary = self._collector()
+            body = self._exporter.export(summary).encode("utf-8")
+        except Exception as exc:  # NOSONAR
+            _log.error("metrics handler error: %s", exc)
+            self.send_response(500)
+            self.end_headers()
+            self.wfile.write(b"Internal Server Error\n")
+            return
+
+        self.send_response(200)
+        self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def log_message(self, fmt: str, *args: Any) -> None:  # pragma: no cover
+        # Suppress default access log to stderr.
+        pass
+
+
+def serve_metrics(
+    port: int = 9090,
+    *,
+    host: str = "127.0.0.1",
+    collector: Any | None = None,
+    namespace: str = "spanforge",
+) -> http.server.HTTPServer:
+    """Start a background HTTP server exposing ``/metrics`` in Prometheus format.
+
+    The server runs in a daemon thread so it stops automatically when the main
+    process exits.
+
+    Args:
+        port:      TCP port to bind.  Defaults to ``9090``.
+        host:      Interface to bind.  Defaults to ``"127.0.0.1"`` (localhost
+                   only).  Set to ``"0.0.0.0"`` to expose on all interfaces
+                   (ensure firewall rules are in place).
+        collector: Optional callable ``() -> MetricsSummary``.  Defaults to
+                   :func:`_collect_live_summary` which reads from SpanForge
+                   internals.
+        namespace: Metric name prefix (default ``"spanforge"``).
+
+    Returns:
+        The running :class:`http.server.HTTPServer` instance.
+
+    Example::
+
+        serve_metrics(port=9090)
+        # Scrape at http://localhost:9090/metrics
+    """
+    exporter = PrometheusMetricsExporter(namespace=namespace)
+    _collector = collector if collector is not None else _collect_live_summary
+
+    class _Handler(_MetricsHTTPHandler):
+        pass
+
+    _Handler._exporter = exporter  # type: ignore[attr-defined]
+    _Handler._collector = staticmethod(_collector)  # type: ignore[attr-defined]
+
+    server = http.server.HTTPServer((host, port), _Handler)
+    thread = threading.Thread(
+        target=server.serve_forever,
+        name=f"spanforge-metrics-{port}",
+        daemon=True,
+    )
+    thread.start()
+    _log.info("spanforge metrics server listening on http://%s:%d/metrics", host, port)
+    return server

spanforge/migrate.py

@@ -0,0 +1,278 @@
+"""Schema migration utilities for spanforge events.
+
+Provides forward-only migration functions to convert events from older schema
+versions to the current version.  Migrations are idempotent — migrating an
+event that is already at the target version returns it unchanged.
+
+Usage
+-----
+::
+
+    from spanforge.migrate import v1_to_v2, migrate_file
+
+    # Single event
+    v2_event = v1_to_v2(v1_event)
+
+    # Bulk file migration
+    stats = migrate_file("audit.jsonl", output="audit_v2.jsonl")
+    print(f"Migrated {stats.migrated} events")
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+__all__ = [
+    "MigrationStats",
+    "migrate_file",
+    "v1_to_v2",
+]
+
+
+@dataclass(frozen=True)
+class MigrationStats:
+    """Result of a bulk migration operation.
+
+    Attributes:
+        total:              Total events processed.
+        migrated:           Events that were upgraded to a new schema version.
+        skipped:            Events already at the target version (not modified).
+        errors:             Events that could not be parsed or migrated.
+        warnings:           Non-fatal warnings encountered during migration.
+        output_path:        Path where the migrated events were written.
+        transformed_fields: Mapping of field names to the count of events
+                            where that field was transformed.
+    """
+
+    total: int
+    migrated: int
+    skipped: int
+    errors: int
+    warnings: list[str] = field(default_factory=list)
+    output_path: str = ""
+    transformed_fields: dict[str, int] = field(default_factory=dict)
+
+
+def _rehash_md5_to_sha256(checksum: str | None, payload: dict[str, Any]) -> str | None:
+    """If *checksum* starts with ``md5:``, recompute as ``sha256:``."""
+    if checksum and checksum.startswith("md5:"):
+        canonical = json.dumps(
+            payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False
+        ).encode("utf-8")
+        return f"sha256:{hashlib.sha256(canonical).hexdigest()}"
+    return checksum
+
+
+def _coerce_tag_values(tags: Any) -> dict[str, str]:
+    """Ensure all tag values are strings."""
+    if not isinstance(tags, dict):
+        return {}
+    return {str(k): str(v) for k, v in tags.items()}
+
+
+def v1_to_v2(event: Any) -> Any:  # noqa: ANN401
+    """Migrate a single event from schema version 1.0 to 2.0.
+
+    Changes applied:
+    * ``schema_version`` is set to ``"2.0"``.
+    * Missing ``org_id`` is set to ``None`` (was not required in v1).
+    * Missing ``team_id`` is set to ``None``.
+    * Payload key ``model`` is normalised to ``model_id`` if present.
+    * ``tags`` is initialised to an empty dict if missing; all values
+      are coerced to strings.
+    * ``checksum`` is re-hashed from md5 to sha256 if applicable.
+
+    If the event is already at version ``"2.0"`` or later, it is returned
+    unchanged (idempotent).
+
+    Args:
+        event: Either an :class:`~spanforge.event.Event` instance or a plain
+               ``dict`` (as loaded from JSONL).
+
+    Returns:
+        The migrated event (same type as input).
+    """
+    from spanforge.event import Event  # noqa: PLC0415
+
+    if isinstance(event, Event):
+        if event.schema_version == "2.0":
+            return event
+        payload = dict(event.payload)
+        # Normalise model → model_id
+        if "model" in payload and "model_id" not in payload:
+            payload["model_id"] = payload.pop("model")
+        # Re-hash md5 checksum
+        checksum = _rehash_md5_to_sha256(event.checksum, payload)
+        # Coerce tag values to strings
+        tags = _coerce_tag_values(event.tags) if event.tags else {}
+        return Event(
+            schema_version="2.0",
+            event_id=event.event_id,
+            event_type=event.event_type,
+            timestamp=event.timestamp,
+            source=event.source,
+            payload=payload,
+            trace_id=event.trace_id,
+            span_id=event.span_id,
+            parent_span_id=event.parent_span_id,
+            org_id=event.org_id,
+            team_id=event.team_id,
+            actor_id=event.actor_id,
+            session_id=event.session_id,
+            tags=tags,
+            checksum=checksum,
+            signature=event.signature,
+            prev_id=event.prev_id,
+        )
+
+    # Dict-based migration  (e.g. raw JSONL parsing)
+    if isinstance(event, dict):
+        if event.get("schema_version") == "2.0":
+            return event
+        d = dict(event)
+        d["schema_version"] = "2.0"
+        d.setdefault("org_id", None)
+        d.setdefault("team_id", None)
+        # Coerce tag values
+        raw_tags = d.get("tags")
+        if isinstance(raw_tags, dict):
+            d["tags"] = {str(k): str(v) for k, v in raw_tags.items()}
+        else:
+            d["tags"] = {}
+        payload = d.get("payload", {})
+        if isinstance(payload, dict):
+            if "model" in payload and "model_id" not in payload:
+                payload["model_id"] = payload.pop("model")
+        # Re-hash md5 checksum
+        if d.get("checksum", "").startswith("md5:") and isinstance(payload, dict):
+            canonical = json.dumps(
+                payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False
+            ).encode("utf-8")
+            d["checksum"] = f"sha256:{hashlib.sha256(canonical).hexdigest()}"
+        return d
+
+    raise TypeError(f"Cannot migrate object of type {type(event).__name__}")
+
+
+def migrate_file(
+    input_path: str | Path,
+    *,
+    output: str | Path | None = None,
+    org_secret: str | None = None,
+    target_version: str = "2.0",
+    dry_run: bool = False,
+) -> MigrationStats:
+    """Migrate all events in a JSONL file from v1 to v2.
+
+    Reads line-by-line, applies :func:`v1_to_v2` to each JSON object, and
+    writes the result to *output* (defaults to ``<input>_v2.jsonl``).
+
+    Args:
+        input_path:     Path to the source JSONL file.
+        output:         Output file path (default: ``<stem>_v2.jsonl``).
+        org_secret:     When provided, re-signs the migrated chain using HMAC.
+        target_version: Target schema version (default ``"2.0"``).
+        dry_run:        When ``True``, report stats without writing output.
+
+    Returns:
+        A :class:`MigrationStats` summarising the operation.
+    """
+    src = Path(input_path)
+    if output is None:
+        dst = src.with_name(f"{src.stem}_v2{src.suffix}")
+    else:
+        dst = Path(output)
+
+    total = 0
+    migrated = 0
+    skipped = 0
+    errors = 0
+    warnings: list[str] = []
+    transformed_fields: dict[str, int] = {}
+
+    migrated_dicts: list[str] = []
+
+    with src.open("r", encoding="utf-8") as fin:
+        for line_no, line in enumerate(fin, 1):
+            line = line.strip()
+            if not line:
+                continue
+            total += 1
+            try:
+                data = json.loads(line)
+            except json.JSONDecodeError:
+                errors += 1
+                migrated_dicts.append(line + "\n")
+                continue
+
+            # Source format validation
+            if not isinstance(data, dict):
+                errors += 1
+                warnings.append(f"line {line_no}: not a JSON object")
+                migrated_dicts.append(line + "\n")
+                continue
+
+            if data.get("schema_version") == target_version:
+                skipped += 1
+                migrated_dicts.append(line + "\n")
+                continue
+
+            try:
+                # Track which fields get transformed
+                payload = data.get("payload", {})
+                if isinstance(payload, dict) and "model" in payload and "model_id" not in payload:
+                    transformed_fields["payload.model→model_id"] = transformed_fields.get("payload.model→model_id", 0) + 1
+                if data.get("checksum", "").startswith("md5:"):
+                    transformed_fields["checksum.md5→sha256"] = transformed_fields.get("checksum.md5→sha256", 0) + 1
+                raw_tags = data.get("tags", {})
+                if isinstance(raw_tags, dict) and any(not isinstance(v, str) for v in raw_tags.values()):
+                    transformed_fields["tags.value_coercion"] = transformed_fields.get("tags.value_coercion", 0) + 1
+
+                migrated_data = v1_to_v2(data)
+                migrated_dicts.append(
+                    json.dumps(migrated_data, separators=(",", ":"), ensure_ascii=False) + "\n"
+                )
+                migrated += 1
+            except Exception:  # NOSONAR
+                errors += 1
+                migrated_dicts.append(line + "\n")
+
+    # Re-sign if org_secret provided
+    if org_secret and not dry_run:
+        from spanforge.event import Event  # noqa: PLC0415
+        from spanforge.signing import sign as _sign  # noqa: PLC0415
+
+        signed_lines: list[str] = []
+        prev_event = None
+        for raw_line in migrated_dicts:
+            raw_line = raw_line.strip()
+            if not raw_line:
+                continue
+            try:
+                data = json.loads(raw_line)
+                evt = Event.from_dict(data)
+                signed_evt = _sign(evt, org_secret, prev_event=prev_event)
+                prev_event = signed_evt
+                signed_lines.append(signed_evt.to_json() + "\n")
+            except Exception:  # noqa: BLE001
+                signed_lines.append(raw_line + "\n")
+        migrated_dicts = signed_lines
+
+    if not dry_run:
+        with dst.open("w", encoding="utf-8") as fout:
+            for out_line in migrated_dicts:
+                fout.write(out_line)
+
+    return MigrationStats(
+        total=total,
+        migrated=migrated,
+        skipped=skipped,
+        errors=errors,
+        warnings=warnings,
+        output_path=str(dst),
+        transformed_fields=transformed_fields,
+    )