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