"""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())