DocForge2/src/docforge/telemetry.py

"""Bounded request-local diagnostics for repository gates and explicit profiling."""

from __future__ import annotations

import time
from collections.abc import Generator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Literal

CounterName = Literal[
    "project_loads",
    "source_files_parsed",
    "source_bytes_parsed",
    "adapter_projection_loads",
    "adapter_source_extractions",
    "source_generation_checks",
    "index_checks",
    "index_synchronizations",
    "index_builds",
    "render_prepare_calls",
    "render_output_bytes_built",
    "render_output_bytes_hashed",
    "viewer_manager_requests",
]
StageName = Literal[
    "source.generation",
    "source.parse",
    "adapter.projection",
    "adapter.extract",
    "index.check",
    "index.synchronize",
    "index.build",
    "index.read",
    "render.status",
    "render.prepare",
    "render.output_hash",
    "visualization.status",
    "viewer.manager",
    "mcp.runtime_validation",
]

COUNTER_NAMES: tuple[CounterName, ...] = (
    "project_loads",
    "source_files_parsed",
    "source_bytes_parsed",
    "adapter_projection_loads",
    "adapter_source_extractions",
    "source_generation_checks",
    "index_checks",
    "index_synchronizations",
    "index_builds",
    "render_prepare_calls",
    "render_output_bytes_built",
    "render_output_bytes_hashed",
    "viewer_manager_requests",
)
STAGE_NAMES: frozenset[StageName] = frozenset(
    {
        "source.generation",
        "source.parse",
        "adapter.projection",
        "adapter.extract",
        "index.check",
        "index.synchronize",
        "index.build",
        "index.read",
        "render.status",
        "render.prepare",
        "render.output_hash",
        "visualization.status",
        "viewer.manager",
        "mcp.runtime_validation",
    }
)
OPERATION_NAMES = frozenset(
    {
        "test",
        "benchmark.m1",
        "benchmark.m2",
        "benchmark.m3",
        "mcp.invoke",
        "mcp.bootstrap",
        "mcp.sync",
        "mcp.project_info",
        "mcp.contract",
        "mcp.get_node",
        "mcp.get_logic",
        "mcp.search",
        "mcp.filter",
        "mcp.backlinks",
        "mcp.dependencies",
        "mcp.impact",
        "mcp.context",
        "mcp.task_context",
        "mcp.generation_diff",
        "mcp.validate_project",
        "mcp.render_status",
        "mcp.graph_plan",
        "mcp.graph_render_status",
        "mcp.visualize",
        "mcp.visualization_status",
        "mcp.stop_visualization",
        "mcp.changeset",
        "mcp.mutation",
        "cli.onboard",
        "cli.info",
        "cli.validate",
        "cli.build",
        "cli.reindex",
        "cli.sync",
        "cli.check",
        "cli.validate-index",
        "cli.show",
        "cli.search",
        "cli.filter",
        "cli.backlinks",
        "cli.dependencies",
        "cli.impact",
        "cli.context",
        "cli.generation-diff",
        "cli.graph-plan",
        "cli.graph-render",
        "cli.graph-render-status",
        "cli.configure",
        "cli.doctor",
        "cli.render",
        "cli.render-status",
        "cli.preview",
        "cli.apply",
        "cli.visualize",
        "cli.visualization-status",
        "cli.visualization-stop",
    }
)


@dataclass
class _StageAggregate:
    calls: int = 0
    elapsed_ns: int = 0


@dataclass
class Collector:
    """One bounded aggregate owned by the current request context."""

    operation: str
    counters: dict[CounterName, int] = field(
        default_factory=lambda: {name: 0 for name in COUNTER_NAMES}
    )
    stages: dict[StageName, _StageAggregate] = field(default_factory=lambda: {})
    elapsed_ns: int = 0

    def as_dict(self, *, outcome: str) -> dict[str, object]:
        if outcome not in {"ok", "error"}:
            raise ValueError("Telemetry outcome must be ok or error")
        return {
            "schema_version": 1,
            "operation": self.operation,
            "outcome": outcome,
            "elapsed_ns": self.elapsed_ns,
            "stages": {
                name: {
                    "calls": aggregate.calls,
                    "elapsed_ns": aggregate.elapsed_ns,
                }
                for name, aggregate in sorted(self.stages.items())
            },
            "counters": {name: self.counters[name] for name in COUNTER_NAMES},
        }


_CURRENT: ContextVar[Collector | None] = ContextVar(
    "docforge_telemetry",
    default=None,
)


@contextmanager
def request(
    operation: str,
    *,
    enabled: bool,
) -> Generator[Collector | None, None, None]:
    """Collect one explicit request without affecting the disabled path."""

    if operation not in OPERATION_NAMES:
        raise ValueError("Unknown telemetry operation")
    if not enabled:
        yield None
        return
    collector = Collector(operation=operation)
    token = _CURRENT.set(collector)
    started = time.perf_counter_ns()
    try:
        yield collector
    finally:
        collector.elapsed_ns = time.perf_counter_ns() - started
        _CURRENT.reset(token)


def increment(counter: CounterName | str, amount: int = 1) -> None:
    """Increment one fixed counter when a request collector is active."""

    if counter not in COUNTER_NAMES:
        raise ValueError("Unknown telemetry counter")
    if type(amount) is not int or amount < 0:
        raise ValueError("Telemetry increments must be nonnegative integers")
    collector = _CURRENT.get()
    if collector is not None:
        collector.counters[counter] += amount


@contextmanager
def stage(name: StageName | str) -> Generator[None, None, None]:
    """Aggregate one fixed stage while avoiding a clock read when disabled."""

    if name not in STAGE_NAMES:
        raise ValueError("Unknown telemetry stage")
    collector = _CURRENT.get()
    if collector is None:
        yield
        return
    started = time.perf_counter_ns()
    try:
        yield
    finally:
        aggregate = collector.stages.setdefault(name, _StageAggregate())
        aggregate.calls += 1
        aggregate.elapsed_ns += time.perf_counter_ns() - started