DocForge2/tools/milestone2_benchmark.py

"""Milestone 2 agent-retrieval and client-integration benchmark gates."""

from __future__ import annotations

import argparse
import hashlib
import json
import platform
import resource
import subprocess
import sys
import tempfile
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.client_config import generate_client_configuration
from docforge.doctor import run_doctor
from docforge.index import ProjectIndex
from docforge.mcp_server import DocForgeService
from docforge.pagination import canonical_hash
from docforge.project import Project
from docforge.retrieval import MAX_TASK_EVIDENCE, build_retrieval_plan
from docforge.telemetry import COUNTER_NAMES, request

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


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Gate DocForge2 Milestone 2 on disposable agent workflows."
    )
    parser.add_argument("--nodes", type=int, default=1000)
    parser.add_argument("--samples", type=int, default=10)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--memory-probe", action="store_true", help=argparse.SUPPRESS)
    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 _compact_size(value: object) -> int:
    return len(
        json.dumps(value, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8")
    )


def _prepare_fixture(root: Path, node_count: int) -> None:
    write_synthetic_project(root, node_count)
    descriptor = root / ".docforge" / "project.toml"
    descriptor.write_text(
        descriptor.read_text(encoding="utf-8")
        .replace("max_results = 100", "max_results = 1000")
        .replace("max_tool_output_chars = 5000000", "max_tool_output_chars = 200000"),
        encoding="utf-8",
    )
    focus = synthetic_node_id(0)
    for index in range(1, node_count):
        path = root / "docs" / "content" / f"node-{index:04d}.md"
        raw = path.read_text(encoding="utf-8")
        previous = synthetic_node_id(index - 1)
        path.write_text(
            raw.replace(
                f'depends_on = ["{previous}"]',
                f'depends_on = ["{focus}"]',
            ),
            encoding="utf-8",
        )


def _tree_hash(root: Path) -> str:
    digest = hashlib.sha256()
    for path in sorted(root.rglob("*")):
        if not path.is_file():
            continue
        digest.update(path.relative_to(root).as_posix().encode("utf-8"))
        digest.update(b"\0")
        digest.update(path.read_bytes())
        digest.update(b"\0")
    return digest.hexdigest()


def _assert_zero_work(diagnostics: Mapping[str, object]) -> None:
    counters_value = diagnostics.get("counters")
    if not isinstance(counters_value, Mapping):
        raise RuntimeError("Measured operation did not expose telemetry counters")
    counters = cast(Mapping[str, object], counters_value)
    for counter in ZERO_WORK_COUNTERS:
        if counters.get(counter) != 0:
            raise RuntimeError(f"Milestone 2 operation performed forbidden work: {counter}")


def _result_diagnostics(
    result: Mapping[str, object],
    *,
    expected_counters: Mapping[str, int],
) -> Mapping[str, object]:
    diagnostics_value = result.get("diagnostics")
    if not isinstance(diagnostics_value, Mapping):
        raise RuntimeError("Measured MCP result did not include diagnostics")
    diagnostics = cast(Mapping[str, object], diagnostics_value)
    counters = cast(Mapping[str, object], diagnostics["counters"])
    for counter in ZERO_WORK_COUNTERS:
        expected = expected_counters.get(counter, 0)
        if counters.get(counter) != expected:
            raise RuntimeError(
                f"Measured MCP result expected {counter}={expected}, "
                f"received {counters.get(counter)!r}"
            )
    for counter, expected in expected_counters.items():
        if counters.get(counter) != expected:
            raise RuntimeError(
                f"Measured MCP result expected {counter}={expected}, "
                f"received {counters.get(counter)!r}"
            )
    return diagnostics


def _maximum_page_validator(
    result: Mapping[str, object],
    *,
    expected_counters: Mapping[str, int],
    diagnostics_dropped: list[bool],
) -> None:
    if isinstance(result.get("diagnostics"), Mapping):
        _result_diagnostics(result, expected_counters=expected_counters)
        diagnostics_dropped.append(False)
        return
    if _compact_size(result) <= 190_000:
        raise RuntimeError(
            "Maximum-page diagnostics were absent before the primary result approached "
            "the response budget"
        )
    diagnostics_dropped.append(True)


def _measure(
    operation: Callable[[], dict[str, object]],
    *,
    samples: int,
    p95_limit_ms: float,
    response_limit_bytes: int,
    validator: Callable[[dict[str, object]], object] | None = None,
) -> tuple[dict[str, object], dict[str, object]]:
    results: list[dict[str, object]] = []
    diagnostics_records: list[Mapping[str, object]] = []

    def validated_operation() -> dict[str, object]:
        result = operation()
        if result.get("status") != "ok":
            raise RuntimeError("Measured operation did not succeed")
        response_bytes = _compact_size(result)
        if response_bytes > response_limit_bytes:
            raise RuntimeError(
                f"Milestone 2 response {response_bytes} exceeds {response_limit_bytes} bytes"
            )
        if validator is not None:
            validator(result)
        diagnostics_value = result.get("diagnostics", result.get("_benchmark_diagnostics"))
        if isinstance(diagnostics_value, Mapping):
            diagnostics_records.append(cast(Mapping[str, object], diagnostics_value))
        results.append(result)
        return result

    measurement, last_value = measure_operation(validated_operation, samples=samples)
    if not isinstance(last_value, Mapping):
        raise RuntimeError("Measured operation returned a non-object result")
    last = dict(cast(Mapping[str, object], last_value))
    p95 = cast(float, measurement["p95_ms"])
    if p95 > p95_limit_ms:
        raise RuntimeError(f"Milestone 2 operation p95 {p95:.3f} ms exceeds {p95_limit_ms:.3f} ms")
    counter_ranges: dict[str, dict[str, int]] = {}
    if diagnostics_records:
        for counter in COUNTER_NAMES:
            values = [
                cast(
                    int,
                    cast(Mapping[str, object], record["counters"])[counter],
                )
                for record in diagnostics_records
            ]
            counter_ranges[counter] = {
                "minimum": min(values),
                "maximum": max(values),
            }
    return (
        {
            **measurement,
            "p95_limit_ms": p95_limit_ms,
            "response_limit_bytes": response_limit_bytes,
            "validated_invocations": len(results),
            "maximum_response_bytes": max(_compact_size(result) for result in results),
            **({"counter_ranges": counter_ranges} if counter_ranges else {}),
        },
        last,
    )


def _profiled(
    operation: Callable[[], dict[str, object]],
) -> dict[str, object]:
    with request("benchmark.m2", enabled=True) as collector:
        result = operation()
    if collector is None:
        raise RuntimeError("Milestone 2 profiling collector was not created")
    diagnostics = collector.as_dict(outcome="ok")
    _assert_zero_work(diagnostics)
    result["_benchmark_diagnostics"] = diagnostics
    return result


def _page_summary(
    pages: list[dict[str, object]],
    *,
    started_ns: int,
) -> dict[str, object]:
    sizes = [_compact_size(page) for page in pages]
    paginations = [cast(dict[str, object], page["pagination"]) for page in pages]
    result: dict[str, object] = {
        "status": "ok",
        "page_count": len(pages),
        "maximum_page_bytes": max(sizes),
        "aggregate_page_bytes": sum(sizes),
        "elapsed_ms": round((time.perf_counter_ns() - started_ns) / 1_000_000, 3),
        "maximum_cursor_bytes": max(
            (
                len(cast(str, page["next_cursor"]).encode("utf-8"))
                for page in paginations
                if page["next_cursor"] is not None
            ),
            default=0,
        ),
    }
    if cast(int, result["maximum_cursor_bytes"]) > 4_096:
        raise RuntimeError("Milestone 2 cursor exceeded 4,096 bytes")
    return result


def _task_oracle(
    service: DocForgeService,
    focus: str,
) -> dict[str, object]:
    plan = build_retrieval_plan(
        service.project.descriptor,
        task_kind="change",
        task="Change the central synthetic workflow",
        focus_node_id=focus,
        budget=None,
        limit=min(
            service.project.descriptor.limits.max_results,
            MAX_TASK_EVIDENCE,
        ),
        effective_policy=service.policy.as_dict(),
    )
    result = service.index.task_context(plan)
    capsule = cast(dict[str, object], result["capsule"])
    generation = cast(Mapping[str, object], capsule["generation"])
    plan_payload = cast(Mapping[str, object], capsule["plan"])
    evidence = cast(list[dict[str, object]], capsule["evidence"])
    gaps = cast(list[dict[str, object]], capsule["gaps"])
    omissions = cast(list[dict[str, object]], capsule["omissions"])
    evidence_hashes = [cast(str, item["evidence_hash"]) for item in evidence]
    collection_hash = canonical_hash(
        {
            "generation": dict(generation),
            "plan_hash": plan_payload["plan_hash"],
            "evidence": evidence_hashes,
            "gaps": gaps,
            "omissions": omissions,
        }
    )
    if capsule["collection_hash"] != collection_hash:
        raise RuntimeError("Task-context full oracle has an invalid collection hash")
    return {
        "generation": dict(generation),
        "plan_hash": plan_payload["plan_hash"],
        "evidence": evidence,
        "evidence_hashes": evidence_hashes,
        "gaps": gaps,
        "omissions": omissions,
        "collection_hash": collection_hash,
    }


def _task_traversal(
    service: DocForgeService,
    focus: str,
    *,
    node_count: int,
    oracle: Mapping[str, object],
) -> dict[str, object]:
    started = time.perf_counter_ns()
    oracle_generation = cast(Mapping[str, object], oracle["generation"])
    oracle_plan_hash = cast(str, oracle["plan_hash"])
    oracle_evidence = cast(list[dict[str, object]], oracle["evidence"])
    oracle_evidence_hashes = cast(list[str], oracle["evidence_hashes"])
    oracle_gaps = cast(list[dict[str, object]], oracle["gaps"])
    oracle_omissions = cast(list[dict[str, object]], oracle["omissions"])
    oracle_collection_hash = cast(str, oracle["collection_hash"])
    oracle_evidence_by_node = {cast(str, item["node_id"]): item for item in oracle_evidence}
    pages: list[dict[str, object]] = []
    cursor: str | None = None
    capsule_hash: str | None = None
    collection_hash: str | None = None
    plan_hash: str | None = None
    returned = 0
    seen_cursors: set[str] = set()
    evidence_hashes: list[str] = []
    evidence_nodes: list[str] = []
    page_omissions: list[dict[str, object]] = []
    omitted_nodes: list[str] = []
    reconstructed_evidence_hashes: list[str] = []
    reconstructed_omissions: list[dict[str, object]] = []
    diagnostics_records: list[Mapping[str, object]] = []
    first_generation: Mapping[str, object] | None = None
    first_gaps: list[object] | None = None
    while True:
        page = service.task_context(
            "change",
            "Change the central synthetic workflow",
            focus_node_id=focus,
            limit=100,
            cursor=cursor,
        )
        if page.get("status") != "ok":
            raise RuntimeError("Task-context page failed")
        diagnostics_records.append(
            _result_diagnostics(
                page,
                expected_counters={
                    "index_checks": 1,
                    "index_synchronizations": 0,
                    "source_generation_checks": 2,
                },
            )
        )
        if _compact_size(page) > 200_000:
            raise RuntimeError("Task-context page exceeded the configured response budget")
        capsule = cast(dict[str, object], page["capsule"])
        plan = cast(dict[str, object], capsule["plan"])
        current_capsule_hash = cast(str, capsule["capsule_hash"])
        current_collection_hash = cast(str, capsule["collection_hash"])
        current_plan_hash = cast(str, plan["plan_hash"])
        generation = cast(Mapping[str, object], capsule["generation"])
        gaps = cast(list[object], capsule["gaps"])
        if first_generation is None:
            first_generation = generation
            first_gaps = gaps
        elif generation != first_generation or gaps != first_gaps:
            raise RuntimeError("Task-context invariant payload changed during traversal")
        if capsule_hash is None:
            capsule_hash = current_capsule_hash
            collection_hash = current_collection_hash
            plan_hash = current_plan_hash
        elif (
            current_capsule_hash != capsule_hash
            or current_collection_hash != collection_hash
            or current_plan_hash != plan_hash
        ):
            raise RuntimeError("Task-context page binding changed during traversal")
        pagination = cast(dict[str, object], page["pagination"])
        current_evidence = cast(list[dict[str, object]], capsule["evidence"])
        current_omissions = cast(list[dict[str, object]], capsule["omissions"])
        if pagination["returned_count"] != len(current_evidence) + len(current_omissions):
            raise RuntimeError("Task-context page count does not match its items")
        for item in current_evidence:
            node_id = cast(str, item["node_id"])
            if oracle_evidence_by_node.get(node_id) != item:
                raise RuntimeError("Task-context page evidence drifted from the full oracle")
            evidence_hash = cast(str, item["evidence_hash"])
            evidence_hashes.append(evidence_hash)
            reconstructed_evidence_hashes.append(evidence_hash)
            evidence_nodes.append(node_id)
        for omission in current_omissions:
            subject = omission.get("subject")
            if not isinstance(subject, str) or omission.get("code") not in {
                "response_limit",
                "token_budget",
            }:
                raise RuntimeError("Task-context omission has unexpected semantics")
            if omission["code"] == "response_limit":
                replaced = oracle_evidence_by_node.get(subject)
                if replaced is None or omission.get("detail_hash") != canonical_hash(replaced):
                    raise RuntimeError(
                        "Task-context response-limit omission does not attest its oracle item"
                    )
                reconstructed_evidence_hashes.append(cast(str, replaced["evidence_hash"]))
            else:
                reconstructed_omissions.append(omission)
            omitted_nodes.append(subject)
        page_omissions.extend(current_omissions)
        returned += cast(int, pagination["returned_count"])
        pages.append(page)
        next_cursor = pagination["next_cursor"]
        if next_cursor is None:
            if returned != pagination["total_count"]:
                raise RuntimeError("Task-context traversal did not reconstruct every item")
            break
        cursor = cast(str, next_cursor)
        if cursor in seen_cursors:
            raise RuntimeError("Task-context pagination repeated a cursor")
        seen_cursors.add(cursor)
    if len(evidence_hashes) != len(set(evidence_hashes)):
        raise RuntimeError("Task-context traversal repeated evidence")
    if len(evidence_nodes) != len(set(evidence_nodes)):
        raise RuntimeError("Task-context traversal repeated node evidence")
    if len(omitted_nodes) != len(set(omitted_nodes)):
        raise RuntimeError("Task-context traversal repeated an omitted subject")
    expected_nodes = {synthetic_node_id(index) for index in range(node_count)}
    reconstructed_nodes = evidence_nodes + omitted_nodes
    if (
        returned != node_count
        or len(reconstructed_nodes) != node_count
        or len(reconstructed_nodes) != len(set(reconstructed_nodes))
        or set(reconstructed_nodes) != expected_nodes
    ):
        raise RuntimeError("Task-context traversal did not reconstruct every synthetic candidate")
    assert first_generation is not None
    assert first_gaps is not None
    if (
        collection_hash != oracle_collection_hash
        or dict(first_generation) != dict(oracle_generation)
        or plan_hash != oracle_plan_hash
        or first_gaps != oracle_gaps
        or reconstructed_evidence_hashes != oracle_evidence_hashes
        or reconstructed_omissions != oracle_omissions
        or collection_hash
        != canonical_hash(
            {
                "generation": dict(first_generation),
                "plan_hash": plan_hash,
                "evidence": reconstructed_evidence_hashes,
                "gaps": first_gaps,
                "omissions": reconstructed_omissions,
            }
        )
    ):
        raise RuntimeError("Task-context traversal did not reconstruct its full-oracle binding")
    summary = _page_summary(pages, started_ns=started)
    if cast(float, summary["elapsed_ms"]) > 2_500:
        raise RuntimeError("Complete task-context traversal exceeded 2,500 ms")
    summary.update(
        {
            "capsule_hash": capsule_hash,
            "collection_hash": collection_hash,
            "plan_hash": plan_hash,
            "item_count": returned,
            "evidence_count": len(evidence_nodes),
            "omission_count": len(omitted_nodes),
            "ordered_evidence_hash": canonical_hash(reconstructed_evidence_hashes),
            "ordered_candidate_hash": canonical_hash(reconstructed_nodes),
            "collection_hash_reconstructed": True,
            "counter_ranges": {
                counter: {
                    "minimum": min(
                        cast(int, cast(Mapping[str, object], item["counters"])[counter])
                        for item in diagnostics_records
                    ),
                    "maximum": max(
                        cast(int, cast(Mapping[str, object], item["counters"])[counter])
                        for item in diagnostics_records
                    ),
                }
                for counter in COUNTER_NAMES
            },
        }
    )
    return summary


def _generation_traversal(
    service: DocForgeService,
    *,
    node_count: int,
) -> dict[str, object]:
    started = time.perf_counter_ns()
    pages: list[dict[str, object]] = []
    cursor: str | None = None
    receipt_hash: str | None = None
    returned = 0
    seen_cursors: set[str] = set()
    item_hashes: list[str] = []
    node_ids: list[str] = []
    diagnostics_records: list[Mapping[str, object]] = []
    retained_collection_hash: str | None = None
    while True:
        page = service.generation_diff(limit=100, cursor=cursor)
        if page.get("status") != "ok":
            raise RuntimeError("Generation-diff page failed")
        diagnostics_records.append(
            _result_diagnostics(
                page,
                expected_counters={
                    "index_checks": 0,
                    "index_synchronizations": 0,
                    "source_generation_checks": 2,
                },
            )
        )
        if _compact_size(page) > 200_000:
            raise RuntimeError("Generation-diff page exceeded the configured response budget")
        generation_diff = cast(dict[str, object], page["generation_diff"])
        header = cast(dict[str, object], generation_diff["receipt_header"])
        current_receipt_hash = cast(str, header["stored_receipt_hash"])
        current_retained_hash = cast(str, header["retained_collection_hash"])
        if receipt_hash is None:
            receipt_hash = current_receipt_hash
            retained_collection_hash = current_retained_hash
        elif current_receipt_hash != receipt_hash:
            raise RuntimeError("Generation-diff receipt changed during traversal")
        elif current_retained_hash != retained_collection_hash:
            raise RuntimeError("Generation-diff collection changed during traversal")
        pagination = cast(dict[str, object], page["pagination"])
        items = cast(list[dict[str, object]], generation_diff["items"])
        omissions = cast(list[dict[str, object]], generation_diff["omissions"])
        if pagination["returned_count"] != len(items) + len(omissions):
            raise RuntimeError("Generation-diff page count does not match its items")
        if omissions:
            raise RuntimeError("Generation-diff traversal omitted a retained item")
        for item in items:
            item_hashes.append(cast(str, item["item_hash"]))
            node_id = item.get("node_id")
            if (
                item.get("entity") != "node"
                or item.get("change") != "changed"
                or not isinstance(node_id, str)
            ):
                raise RuntimeError("Generation-diff synthetic item has unexpected semantics")
            node_ids.append(node_id)
        returned += cast(int, pagination["returned_count"])
        pages.append(page)
        next_cursor = pagination["next_cursor"]
        if next_cursor is None:
            if returned != pagination["total_count"]:
                raise RuntimeError("Generation-diff traversal did not reconstruct every item")
            break
        cursor = cast(str, next_cursor)
        if cursor in seen_cursors:
            raise RuntimeError("Generation-diff pagination repeated a cursor")
        seen_cursors.add(cursor)
    if len(item_hashes) != len(set(item_hashes)):
        raise RuntimeError("Generation-diff traversal repeated a retained item")
    if (
        returned != node_count
        or len(item_hashes) != node_count
        or len(node_ids) != node_count
        or set(node_ids) != {synthetic_node_id(index) for index in range(node_count)}
    ):
        raise RuntimeError("Generation-diff traversal did not reconstruct every changed node")
    if retained_collection_hash != canonical_hash(item_hashes):
        raise RuntimeError("Generation-diff traversal did not reconstruct its collection hash")
    summary = _page_summary(pages, started_ns=started)
    if cast(float, summary["elapsed_ms"]) > 500:
        raise RuntimeError("Complete generation-diff traversal exceeded 500 ms")
    summary.update(
        {
            "receipt_hash": receipt_hash,
            "item_count": returned,
            "ordered_item_hash": canonical_hash(item_hashes),
            "counter_ranges": {
                counter: {
                    "minimum": min(
                        cast(int, cast(Mapping[str, object], item["counters"])[counter])
                        for item in diagnostics_records
                    ),
                    "maximum": max(
                        cast(int, cast(Mapping[str, object], item["counters"])[counter])
                        for item in diagnostics_records
                    ),
                }
                for counter in COUNTER_NAMES
            },
        }
    )
    return summary


def _benchmark(root: Path, node_count: int, samples: int) -> dict[str, object]:
    project = Project.open(root)
    index = ProjectIndex(project)
    initial_snapshot = project.load()
    focus = synthetic_node_id(0)
    if len(initial_snapshot.edges) != node_count - 1 or any(
        edge.target_id != focus or edge.source_id == focus for edge in initial_snapshot.edges
    ):
        raise RuntimeError("Milestone 2 fixture is not the expected focus fan-in graph")
    index.build()
    service = DocForgeService(project, capability_mode_name="read", diagnostics=True)
    no_ast_service = DocForgeService(
        project,
        capability_mode_name="read",
        no_ast=True,
        diagnostics=True,
    )
    operations: dict[str, object] = {}

    operations["bootstrap_read"], bootstrap = _measure(
        service.bootstrap,
        samples=samples,
        p95_limit_ms=100,
        response_limit_bytes=32_768,
        validator=lambda result: _result_diagnostics(
            result,
            expected_counters={
                "index_checks": 1,
                "index_synchronizations": 1,
                "source_generation_checks": 1,
            },
        ),
    )
    bootstrap_policy = cast(dict[str, object], bootstrap["effective_policy"])
    if bootstrap_policy["capability_mode"] != "read":
        raise RuntimeError("Read bootstrap did not preserve the explicit capability mode")
    operations["bootstrap_no_ast"], no_ast_bootstrap = _measure(
        no_ast_service.bootstrap,
        samples=samples,
        p95_limit_ms=100,
        response_limit_bytes=32_768,
        validator=lambda result: _result_diagnostics(
            result,
            expected_counters={
                "index_checks": 1,
                "index_synchronizations": 1,
                "source_generation_checks": 1,
            },
        ),
    )
    no_ast_policy = cast(dict[str, object], no_ast_bootstrap["adapter_policy"])
    if no_ast_policy["mode"] != "preserve-no-ast":
        raise RuntimeError("No-AST bootstrap did not preserve the adapter policy")

    task_oracle = _task_oracle(service, focus)
    task_probe = service.task_context(
        "change",
        "Change the central synthetic workflow",
        focus_node_id=focus,
        limit=1,
    )
    _result_diagnostics(
        task_probe,
        expected_counters={
            "index_checks": 1,
            "index_synchronizations": 0,
            "source_generation_checks": 2,
        },
    )
    gap = service.task_context(
        "implementation",
        "Implement the central synthetic workflow",
        focus_node_id=focus,
        limit=1,
    )
    gap_codes = {
        cast(str, item["code"]) for item in cast(list[dict[str, object]], gap["capsule"]["gaps"])
    }
    if "category_not_declared" not in gap_codes:
        raise RuntimeError("Task-context evidence-gap diagnostic was not preserved")
    operations["task_context_diagnostic_page"], _ = _measure(
        lambda: service.task_context(
            "change",
            "Change the central synthetic workflow",
            focus_node_id=focus,
            limit=100,
        ),
        samples=samples,
        p95_limit_ms=500,
        response_limit_bytes=200_000,
        validator=lambda result: _result_diagnostics(
            result,
            expected_counters={
                "index_checks": 1,
                "index_synchronizations": 0,
                "source_generation_checks": 2,
            },
        ),
    )
    task_diagnostics_dropped: list[bool] = []
    operations["task_context_maximum_page"], _ = _measure(
        lambda: service.task_context(
            "change",
            "Change the central synthetic workflow",
            focus_node_id=focus,
            limit=1_000,
        ),
        samples=samples,
        p95_limit_ms=500,
        response_limit_bytes=200_000,
        validator=lambda result: _maximum_page_validator(
            result,
            expected_counters={
                "index_checks": 1,
                "index_synchronizations": 0,
                "source_generation_checks": 2,
            },
            diagnostics_dropped=task_diagnostics_dropped,
        ),
    )
    operations["task_context_maximum_page"]["diagnostics_dropped_for_budget"] = any(
        task_diagnostics_dropped
    )
    task_order_hash: str | None = None

    def validate_task_summary(result: dict[str, object]) -> None:
        nonlocal task_order_hash
        current = cast(str, result["ordered_evidence_hash"])
        if task_order_hash is None:
            task_order_hash = current
        elif current != task_order_hash:
            raise RuntimeError("Task-context traversal order changed across samples")

    task_measurement, task_summary = _measure(
        lambda: _task_traversal(
            service,
            focus,
            node_count=node_count,
            oracle=task_oracle,
        ),
        samples=samples,
        p95_limit_ms=2_500,
        response_limit_bytes=32_768,
        validator=validate_task_summary,
    )
    operations["task_context_complete"] = {
        **task_measurement,
        "result_summary": task_summary,
    }

    baseline = service.generation_diff(limit=1)
    if baseline.get("receipt_state") != "current":
        raise RuntimeError("Initial generation-diff receipt is not current")
    for path in sorted((root / "docs" / "content").glob("node-*.md")):
        path.write_text(
            path.read_text(encoding="utf-8") + "\nMilestone 2 transition generation.\n",
            encoding="utf-8",
        )
    index.build()
    generation_probe = service.generation_diff(limit=1)
    _result_diagnostics(
        generation_probe,
        expected_counters={
            "index_checks": 0,
            "index_synchronizations": 0,
            "source_generation_checks": 2,
        },
    )
    operations["generation_diff_diagnostic_page"], _ = _measure(
        lambda: service.generation_diff(limit=100),
        samples=samples,
        p95_limit_ms=100,
        response_limit_bytes=200_000,
        validator=lambda result: _result_diagnostics(
            result,
            expected_counters={
                "index_checks": 0,
                "index_synchronizations": 0,
                "source_generation_checks": 2,
            },
        ),
    )
    generation_diagnostics_dropped: list[bool] = []
    operations["generation_diff_maximum_page"], _ = _measure(
        lambda: service.generation_diff(limit=1_000),
        samples=samples,
        p95_limit_ms=100,
        response_limit_bytes=200_000,
        validator=lambda result: _maximum_page_validator(
            result,
            expected_counters={
                "index_checks": 0,
                "index_synchronizations": 0,
                "source_generation_checks": 2,
            },
            diagnostics_dropped=generation_diagnostics_dropped,
        ),
    )
    operations["generation_diff_maximum_page"]["diagnostics_dropped_for_budget"] = any(
        generation_diagnostics_dropped
    )
    generation_order_hash: str | None = None

    def validate_generation_summary(result: dict[str, object]) -> None:
        nonlocal generation_order_hash
        current = cast(str, result["ordered_item_hash"])
        if generation_order_hash is None:
            generation_order_hash = current
        elif current != generation_order_hash:
            raise RuntimeError("Generation-diff traversal order changed across samples")

    generation_measurement, generation_summary = _measure(
        lambda: _generation_traversal(service, node_count=node_count),
        samples=samples,
        p95_limit_ms=500,
        response_limit_bytes=32_768,
        validator=validate_generation_summary,
    )
    operations["generation_diff_complete"] = {
        **generation_measurement,
        "result_summary": generation_summary,
    }

    project_tree_before = _tree_hash(root)
    configurations: dict[str, dict[str, object]] = {}
    for client in ("codex", "claude", "openclaw"):
        configuration_hash: str | None = None

        def validate_configuration(
            result: dict[str, object],
        ) -> None:
            nonlocal configuration_hash
            diagnostics = cast(
                Mapping[str, object],
                result["_benchmark_diagnostics"],
            )
            _assert_zero_work(diagnostics)
            counters = cast(Mapping[str, object], diagnostics["counters"])
            if counters["index_checks"] != 0 or counters["source_generation_checks"] != 0:
                raise RuntimeError("Configuration preview performed hidden project work")
            current = cast(str, result["configuration_hash"])
            if configuration_hash is None:
                configuration_hash = current
            elif current != configuration_hash:
                raise RuntimeError("Configuration preview is not deterministic")

        measurement, result = _measure(
            lambda selected=client: _profiled(
                lambda: generate_client_configuration(project, selected)
            ),
            samples=samples,
            p95_limit_ms=500,
            response_limit_bytes=32_768,
            validator=validate_configuration,
        )
        artifact = cast(dict[str, object], result["artifact"])
        configurations[client] = {
            **measurement,
            "configuration_hash": result["configuration_hash"],
            "artifact_format": artifact["format"],
        }
    operations["configuration_preview"] = configurations
    if _tree_hash(root) != project_tree_before:
        raise RuntimeError("Configuration preview changed the project tree")

    doctors: dict[str, dict[str, object]] = {}
    for client in ("codex", "claude", "openclaw"):
        config_path = root.parent / f"doctor-{client}.config"
        generated = generate_client_configuration(
            project,
            client,
            output=config_path,
        )
        generated_name = cast(str, generated["server_name"])
        config_before = config_path.read_bytes()
        expected_state = "degraded" if client == "claude" else "healthy"

        def validate_doctor(
            result: dict[str, object],
            *,
            selected_client: str = client,
            selected_state: str = expected_state,
        ) -> None:
            diagnostics = cast(
                Mapping[str, object],
                result["_benchmark_diagnostics"],
            )
            _assert_zero_work(diagnostics)
            counters = cast(Mapping[str, object], diagnostics["counters"])
            if counters["index_checks"] != 0 or counters["source_generation_checks"] != 0:
                raise RuntimeError("Doctor performed hidden project work")
            if result["doctor_state"] != selected_state:
                raise RuntimeError(f"Generated {selected_client} configuration did not pass doctor")

        doctor_measurement, doctor_result = _measure(
            lambda selected=client, path=config_path, name=generated_name: _profiled(
                lambda: run_doctor(
                    project,
                    selected,
                    config_path=path,
                    server_name=name,
                )
            ),
            samples=samples,
            p95_limit_ms=100,
            response_limit_bytes=32_768,
            validator=validate_doctor,
        )
        if _tree_hash(root) != project_tree_before or config_path.read_bytes() != config_before:
            raise RuntimeError("Doctor changed project or client configuration state")
        doctors[client] = {
            **doctor_measurement,
            "doctor_state": doctor_result["doctor_state"],
            "summary": doctor_result["summary"],
        }
    operations["doctor"] = doctors
    return {
        "fixture": {
            "kind": "synthetic_generic_focus_fan_in",
            "node_count": node_count,
            "edge_count": node_count - 1,
            "source_file_count": node_count,
            "max_tool_output_chars": 200_000,
        },
        "operations": operations,
        "process_peak_rss_kib": int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss),
    }


def _isolated_memory(nodes: int, samples: int) -> int:
    completed = subprocess.run(
        [
            sys.executable,
            str(Path(__file__).resolve()),
            "--nodes",
            str(nodes),
            "--samples",
            str(samples),
            "--memory-probe",
        ],
        cwd=ROOT,
        check=True,
        capture_output=True,
        text=True,
        timeout=180,
    )
    payload = cast(dict[str, object], json.loads(completed.stdout))
    return cast(int, payload["peak_rss_kib"])


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-milestone2-") as directory:
        root = (Path(directory) / "project").resolve()
        _prepare_fixture(root, arguments.nodes)
        measurement = _benchmark(root, arguments.nodes, arguments.samples)
    if arguments.memory_probe:
        sys.stdout.write(
            json.dumps(
                {"peak_rss_kib": measurement["process_peak_rss_kib"]},
                sort_keys=True,
            )
        )
        return 0
    isolated_peak = _isolated_memory(arguments.nodes, arguments.samples)
    if isolated_peak > 262_144:
        raise RuntimeError("Milestone 2 isolated process exceeded 256 MiB peak RSS")
    result = {
        "schema_version": 1,
        "benchmark": "docforge2_milestone2",
        "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",
            "response_size": "UTF-8 bytes of compact sorted JSON",
            "samples": arguments.samples,
            "memory_probe_samples": arguments.samples,
            "warmups": 1,
            "percentile": "nearest-rank",
            "memory": "isolated child-process resource.getrusage(RUSAGE_SELF).ru_maxrss",
            "memory_limit_kib": 262_144,
            "zero_work_counters": list(ZERO_WORK_COUNTERS),
        },
        **measurement,
        "isolated_process_peak_rss_kib": isolated_peak,
    }
    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())