"""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 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.context import compile_context 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 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 _operation( operation: Callable[[], dict[str, object]], *, samples: int, p95_limit_ms: float, expected_status: str = "ok", ) -> dict[str, object]: measurement, last = measure_operation(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}") diagnostics = _diagnostics(last_payload) 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") return { **measurement, "p95_limit_ms": p95_limit_ms, "diagnostics": diagnostics, } 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) service.visualization = ViewerManagerClient( service.index, state_path=root / ".docforge" / "missing-viewer-manager.json", ) target = synthetic_node_id(node_count - 1) operations = { "warm_no_change_synchronize": _operation( service.synchronize, samples=samples, p95_limit_ms=100, ), "exact_node": _operation( lambda: service.invoke( lambda: service.index.get_node(target), operation_name="mcp.get_node", ), samples=samples, p95_limit_ms=50, ), "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", ), "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, ), "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, ), "context_32k": _operation( lambda: service.invoke( lambda: compile_context(service.index, "active", 32_000), operation_name="mcp.context", ), samples=samples, p95_limit_ms=250, ), "render_receipt_status": _operation( lambda: service.render_status("manual"), samples=samples, p95_limit_ms=50, ), "visualization_unavailable_status": _operation( service.visualization_status, samples=samples, p95_limit_ms=50, expected_status="error", ), } 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", "response_size": "UTF-8 bytes of compact sorted JSON", "samples": arguments.samples, "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())