DocForge2/tools/milestone1_benchmark.py

"""Milestone 1 warm-operation benchmark with algorithmic zero-work gates."""

from __future__ import annotations

import argparse
import json
import platform
import resource
import subprocess
import sys
import tempfile
import threading
import time
from collections.abc import Callable, Mapping
from pathlib import Path
from typing import cast

from milestone0_baseline import (
    measure_operation,
    synthetic_node_id,
    write_synthetic_project,
)

from docforge.index import ProjectIndex
from docforge.mcp_server import DocForgeService
from docforge.project import Project
from docforge.rendering import RenderService
from docforge.viewer_manager import ViewerManager, ViewerManagerClient

ROOT = Path(__file__).resolve().parents[1]
ZERO_WORK_COUNTERS = (
    "project_loads",
    "source_files_parsed",
    "source_bytes_parsed",
    "adapter_projection_loads",
    "adapter_source_extractions",
    "index_builds",
    "render_prepare_calls",
    "render_output_bytes_built",
    "render_output_bytes_hashed",
)


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Gate warm DocForge2 core work on a disposable deterministic project."
    )
    parser.add_argument("--nodes", type=int, default=1000)
    parser.add_argument("--samples", type=int, default=10)
    parser.add_argument("--output", type=Path)
    return parser


def _git(command: list[str]) -> str:
    return subprocess.run(
        ["git", *command],
        cwd=ROOT,
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()


def _diagnostics(result: object) -> Mapping[str, object]:
    if not isinstance(result, Mapping):
        raise RuntimeError("Measured operation returned a non-object result")
    result_payload = cast(Mapping[str, object], result)
    diagnostics_value = result_payload.get("diagnostics")
    if not isinstance(diagnostics_value, Mapping):
        raise RuntimeError("Measured operation did not return diagnostics")
    diagnostics = cast(Mapping[str, object], diagnostics_value)
    counters_value = diagnostics.get("counters")
    if not isinstance(counters_value, Mapping):
        raise RuntimeError("Measured diagnostics did not return counters")
    counters = cast(Mapping[str, object], counters_value)
    for counter in ZERO_WORK_COUNTERS:
        if counters.get(counter) != 0:
            raise RuntimeError(f"Warm operation performed forbidden work: {counter}")
    return diagnostics


def _result_summary(result: Mapping[str, object]) -> dict[str, object]:
    """Retain bounded semantic evidence without copying primary result payloads."""

    summary: dict[str, object] = {}
    for key in (
        "status",
        "count",
        "limit",
        "truncated",
        "truncation_reason",
        "candidate_edges_consumed",
        "candidate_edges_limit",
        "budget",
        "estimated_tokens",
        "state",
        "verification",
        "configured",
        "snapshot_state",
        "staleness",
    ):
        if key in result:
            summary[key] = result[key]
    error = result.get("error")
    if isinstance(error, Mapping):
        error_payload = cast(Mapping[str, object], error)
        if isinstance(error_payload.get("code"), str):
            summary["error_code"] = error_payload["code"]
    synchronization = result.get("synchronization")
    if isinstance(synchronization, Mapping):
        synchronization_payload = cast(Mapping[str, object], synchronization)
        if isinstance(synchronization_payload.get("action"), str):
            summary["synchronization_action"] = synchronization_payload["action"]
    freshness = result.get("freshness")
    if isinstance(freshness, Mapping):
        freshness_payload = cast(Mapping[str, object], freshness)
        summary["freshness"] = {
            key: freshness_payload[key]
            for key in ("index", "source")
            if isinstance(freshness_payload.get(key), str)
        }
    outputs = result.get("outputs")
    if isinstance(outputs, list):
        summarized_outputs: list[dict[str, object]] = []
        for item in cast(list[object], outputs)[:10]:
            if not isinstance(item, Mapping):
                continue
            item_payload = cast(Mapping[str, object], item)
            summarized_outputs.append(
                {
                    key: item_payload[key]
                    for key in ("view_id", "state", "reason")
                    if key in item_payload
                }
            )
        summary["outputs"] = summarized_outputs
    entries = result.get("entries")
    if isinstance(entries, list):
        summary["entry_count"] = len(cast(list[object], entries))
    omissions = result.get("omissions")
    if isinstance(omissions, list):
        summary["omission_count"] = len(cast(list[object], omissions))
    pagination = result.get("pagination")
    if isinstance(pagination, Mapping):
        pagination_payload = cast(Mapping[str, object], pagination)
        summary["pagination"] = {
            key: pagination_payload[key]
            for key in ("kind", "returned_count", "limit", "total_count", "has_more")
            if key in pagination_payload
        }
    return summary


def _operation(
    operation: Callable[[], dict[str, object]],
    *,
    samples: int,
    p95_limit_ms: float,
    expected_status: str = "ok",
    expected_counters: Mapping[str, int],
) -> dict[str, object]:
    diagnostics_records: list[Mapping[str, object]] = []
    result_summaries: list[dict[str, object]] = []

    def validated_operation() -> dict[str, object]:
        result = operation()
        diagnostics = _diagnostics(result)
        counters_value = diagnostics["counters"]
        if not isinstance(counters_value, Mapping):
            raise RuntimeError("Measured diagnostics did not return counters")
        counters = cast(Mapping[str, object], counters_value)
        for counter, expected in expected_counters.items():
            if counters.get(counter) != expected:
                raise RuntimeError(
                    f"Warm operation expected {counter}={expected}, "
                    f"received {counters.get(counter)!r}"
                )
        diagnostics_records.append(diagnostics)
        result_summaries.append(_result_summary(result))
        return result

    measurement, last = measure_operation(validated_operation, samples=samples)
    if not isinstance(last, Mapping):
        raise RuntimeError(f"Measured operation did not return status={expected_status}")
    last_payload = cast(Mapping[str, object], last)
    if last_payload.get("status") != expected_status:
        raise RuntimeError(f"Measured operation did not return status={expected_status}")
    for summary in result_summaries:
        if summary.get("status") != expected_status:
            raise RuntimeError(f"Measured operation did not return status={expected_status}")
    p95_ms = float(cast(float, measurement["p95_ms"]))
    if p95_ms > p95_limit_ms:
        raise RuntimeError(f"Warm operation p95 {p95_ms:.3f} ms exceeds {p95_limit_ms:.3f} ms")
    counter_names = sorted(
        {
            key
            for diagnostics in diagnostics_records
            for key in cast(Mapping[str, object], diagnostics["counters"])
        }
    )
    counter_ranges = {
        counter: {
            "minimum": min(
                cast(int, cast(Mapping[str, object], record["counters"])[counter])
                for record in diagnostics_records
            ),
            "maximum": max(
                cast(int, cast(Mapping[str, object], record["counters"])[counter])
                for record in diagnostics_records
            ),
        }
        for counter in counter_names
    }
    return {
        **measurement,
        "p95_limit_ms": p95_limit_ms,
        "validated_invocations": len(diagnostics_records),
        "counter_expectations": dict(sorted(expected_counters.items())),
        "counter_ranges": counter_ranges,
        "result_summary": result_summaries[-1],
    }


def _benchmark(root: Path, node_count: int, samples: int) -> dict[str, object]:
    project = Project.open(root)
    ProjectIndex(project).build()
    RenderService(project).render("manual")
    service = DocForgeService(project, diagnostics=True)
    target = synthetic_node_id(node_count - 1)
    backlink_target = synthetic_node_id(node_count - 2)
    first = synthetic_node_id(0)
    read_counters = {
        "index_checks": 1,
        "index_synchronizations": 0,
        "viewer_manager_requests": 0,
    }
    status_counters = {
        "index_checks": 0,
        "index_synchronizations": 0,
        "viewer_manager_requests": 0,
    }
    visualization_counters = {
        "index_checks": 0,
        "index_synchronizations": 0,
        "viewer_manager_requests": 1,
    }
    operations = {
        "warm_no_change_synchronize": _operation(
            service.synchronize,
            samples=samples,
            p95_limit_ms=100,
            expected_counters={
                "index_checks": 1,
                "index_synchronizations": 1,
                "viewer_manager_requests": 0,
            },
        ),
        "exact_node": _operation(
            lambda: service.invoke(
                lambda: service.index.get_node(target),
                operation_name="mcp.get_node",
            ),
            samples=samples,
            p95_limit_ms=50,
            expected_counters=read_counters,
        ),
        "missing_node_error": _operation(
            lambda: service.invoke(
                lambda: service.index.get_node("missing.node"),
                operation_name="mcp.get_node",
            ),
            samples=samples,
            p95_limit_ms=50,
            expected_status="error",
            expected_counters=read_counters,
        ),
        "search_limit_20": _operation(
            lambda: service.invoke(
                lambda: service.index.search("Synthetic measurement", limit=20),
                operation_name="mcp.search",
            ),
            samples=samples,
            p95_limit_ms=100,
            expected_counters=read_counters,
        ),
        "filter_limit_20": _operation(
            lambda: service.invoke(
                lambda: service.index.filter_nodes(family="guide", limit=20),
                operation_name="mcp.filter",
            ),
            samples=samples,
            p95_limit_ms=100,
            expected_counters=read_counters,
        ),
        "backlinks_limit_20": _operation(
            lambda: service.invoke(
                lambda: service.index.backlinks(backlink_target, limit=20),
                operation_name="mcp.backlinks",
            ),
            samples=samples,
            p95_limit_ms=100,
            expected_counters=read_counters,
        ),
        "dependencies_depth_8": _operation(
            lambda: service.invoke(
                lambda: service.index.dependencies(target, depth=8, limit=100),
                operation_name="mcp.dependencies",
            ),
            samples=samples,
            p95_limit_ms=100,
            expected_counters=read_counters,
        ),
        "impact_depth_8": _operation(
            lambda: service.invoke(
                lambda: service.index.impact(first, depth=8, limit=100),
                operation_name="mcp.impact",
            ),
            samples=samples,
            p95_limit_ms=100,
            expected_counters=read_counters,
        ),
        "context_32k": _operation(
            lambda: service.context("active", 32_000, limit=20),
            samples=samples,
            p95_limit_ms=250,
            expected_counters=read_counters,
        ),
        "render_receipt_status": _operation(
            lambda: service.render_status("manual"),
            samples=samples,
            p95_limit_ms=50,
            expected_counters=status_counters,
        ),
    }
    template = root / "docs" / "templates" / "manual.html"
    original_template = template.read_bytes()
    template.write_bytes(original_template + b"\n")
    operations["render_stale_status"] = _operation(
        lambda: service.render_status("manual"),
        samples=samples,
        p95_limit_ms=50,
        expected_counters=status_counters,
    )
    template.write_bytes(original_template)
    RenderService(project).render("manual")
    receipt_path = root / ".docforge" / "cache" / "render-receipts" / "manual.json"
    receipt_path.unlink()
    operations["render_missing_receipt_status"] = _operation(
        lambda: service.render_status("manual"),
        samples=samples,
        p95_limit_ms=50,
        expected_counters=status_counters,
    )
    RenderService(project).render("manual")
    receipt_path.write_text("{", encoding="utf-8")
    operations["render_corrupt_receipt_status"] = _operation(
        lambda: service.render_status("manual"),
        samples=samples,
        p95_limit_ms=50,
        expected_counters=status_counters,
    )
    RenderService(project).render("manual")
    state_path = root / ".docforge" / "benchmark-viewer-manager.json"
    manager = ViewerManager(state_path, check_interval_seconds=0.02)
    manager_thread = threading.Thread(target=manager.serve_forever, daemon=True)
    manager_thread.start()
    deadline = time.monotonic() + 2
    while not state_path.exists() and time.monotonic() < deadline:
        time.sleep(0.01)
    if not state_path.exists():
        manager.shutdown()
        manager_thread.join(timeout=2)
        raise RuntimeError("Viewer manager did not start")
    service.visualization = ViewerManagerClient(service.index, state_path=state_path)
    try:
        service.visualization.start()
        operations["visualization_current_status"] = _operation(
            service.visualization_status,
            samples=samples,
            p95_limit_ms=50,
            expected_counters=visualization_counters,
        )
        with service.index.path.open("ab") as stream:
            stream.write(b"\n")
        operations["visualization_stale_status"] = _operation(
            service.visualization_status,
            samples=samples,
            p95_limit_ms=50,
            expected_counters=visualization_counters,
        )
        service.stop_visualization()
        operations["visualization_not_running_status"] = _operation(
            service.visualization_status,
            samples=samples,
            p95_limit_ms=50,
            expected_counters=visualization_counters,
        )
    finally:
        manager.shutdown()
        manager_thread.join(timeout=2)
    service.visualization = ViewerManagerClient(service.index, state_path=state_path)
    operations["visualization_unavailable_status"] = _operation(
        service.visualization_status,
        samples=samples,
        p95_limit_ms=50,
        expected_status="error",
        expected_counters=visualization_counters,
    )
    return {
        "fixture": {
            "kind": "synthetic_generic",
            "node_count": node_count,
            "edge_count": node_count - 1,
            "source_file_count": node_count,
        },
        "operations": operations,
        "process_peak_rss_kib": int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss),
    }


def main() -> int:
    arguments = _parser().parse_args()
    if arguments.nodes < 2:
        raise SystemExit("--nodes must be at least 2")
    if arguments.samples < 1:
        raise SystemExit("--samples must be positive")
    with tempfile.TemporaryDirectory(prefix="docforge-milestone1-") as directory:
        root = Path(directory).resolve()
        write_synthetic_project(root, arguments.nodes)
        measurement = _benchmark(root, arguments.nodes, arguments.samples)
    result = {
        "schema_version": 1,
        "benchmark": "docforge2_milestone1",
        "source": {
            "revision": _git(["rev-parse", "HEAD"]),
            "dirty": bool(_git(["status", "--porcelain"])),
        },
        "environment": {
            "platform": platform.platform(),
            "machine": platform.machine(),
            "python": platform.python_version(),
            "implementation": platform.python_implementation(),
        },
        "method": {
            "clock": "time.perf_counter_ns",
            "memory": "resource.getrusage(RUSAGE_SELF).ru_maxrss",
            "memory_scope": (
                "cumulative main-process high-water mark; detached viewer-worker memory excluded"
            ),
            "response_size": "UTF-8 bytes of compact sorted JSON",
            "samples": arguments.samples,
            "warmups": 1,
            "percentile": "nearest-rank",
            "zero_work_counters": list(ZERO_WORK_COUNTERS),
        },
        **measurement,
    }
    encoded = json.dumps(result, sort_keys=True, indent=2) + "\n"
    if arguments.output is not None:
        output = arguments.output.resolve()
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(encoded, encoding="utf-8")
    sys.stdout.write(encoded)
    return 0


if __name__ == "__main__":
    sys.exit(main())