cleveragents-core/robot/helper_scaling_performance.py

"""Robot Framework helper for large-project scaling performance tests.

Provides a CLI-style interface for Robot to invoke extended fixture
validation, threshold sub-linearity checks, and a small-scale indexing
throughput test.

Usage:
    python robot/helper_scaling_performance.py check-extended-profiles
    python robot/helper_scaling_performance.py check-extended-thresholds
    python robot/helper_scaling_performance.py check-context-assembly
    python robot/helper_scaling_performance.py check-execution-throughput
    python robot/helper_scaling_performance.py check-sublinear
    python robot/helper_scaling_performance.py run-small-index
"""

from __future__ import annotations

import json
import math
import os
import shutil
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

# Ensure the src directory is on the import path.
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

_FIXTURES_DIR = Path(__file__).resolve().parents[1] / "features" / "fixtures" / "scale"


def _load_json(filename: str) -> dict[str, Any]:
    """Load a JSON fixture file."""
    path = _FIXTURES_DIR / filename
    if not path.exists():
        print(f"ERROR: fixture not found: {path}")
        sys.exit(1)
    with open(path) as fh:
        result: dict[str, Any] = json.load(fh)
        return result


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------


def _cmd_check_extended_profiles() -> int:
    """Verify xlarge and xxlarge profiles are present with valid mixes."""
    data = _load_json("scale_metadata.json")
    names = {p["name"] for p in data["scale_profiles"]}
    assert "xlarge" in names, "Missing xlarge profile"
    assert "xxlarge" in names, "Missing xxlarge profile"
    for profile in data["scale_profiles"]:
        if profile["name"] in ("xlarge", "xxlarge"):
            total = sum(profile["language_mix"].values())
            assert math.isclose(total, 1.0, rel_tol=1e-6), (
                f"Language mix for {profile['name']} sums to {total}"
            )
            assert profile["file_count"] >= 50000, (
                f"File count for {profile['name']} too low: {profile['file_count']}"
            )
    print("scaling-profiles-ok: xlarge and xxlarge profiles validated")
    return 0


def _cmd_check_extended_thresholds() -> int:
    """Verify 50K and 100K indexing and decomposition thresholds."""
    data = _load_json("baseline_thresholds.json")
    for section in ("indexing", "decomposition"):
        for key in ("50000_files", "100000_files"):
            assert key in data[section], f"Missing {key} in {section}"
            t = data[section][key]
            for pct in ("p50_ms", "p95_ms", "p99_ms"):
                assert pct in t, f"Missing {pct} in {section}[{key}]"
            assert t["p50_ms"] < t["p95_ms"] < t["p99_ms"], (
                f"Non-monotonic in {section}[{key}]"
            )
    print("scaling-thresholds-ok: 50K and 100K thresholds validated")
    return 0


def _cmd_check_context_assembly() -> int:
    """Verify context assembly threshold section exists."""
    data = _load_json("baseline_thresholds.json")
    assert "context_assembly" in data, "Missing context_assembly section"
    for key in ("5000_fragments", "10000_fragments"):
        assert key in data["context_assembly"], f"Missing {key} in context_assembly"
    print("scaling-ctx-assembly-ok: context assembly thresholds validated")
    return 0


def _cmd_check_execution_throughput() -> int:
    """Verify execution throughput threshold section exists."""
    data = _load_json("baseline_thresholds.json")
    assert "execution_throughput" in data, "Missing execution_throughput section"
    for key in ("10_plans", "50_plans", "100_plans"):
        assert key in data["execution_throughput"], (
            f"Missing {key} in execution_throughput"
        )
    print("scaling-exec-throughput-ok: execution throughput thresholds validated")
    return 0


def _cmd_check_sublinear() -> int:
    """Verify sub-linear scaling of extended thresholds."""
    data = _load_json("baseline_thresholds.json")
    idx = data["indexing"]
    p50_10k = idx["10000_files"]["p50_ms"]
    p50_50k = idx["50000_files"]["p50_ms"]
    p50_100k = idx["100000_files"]["p50_ms"]
    # 50K should be less than 5x the 10K value (sub-linear)
    assert p50_50k < 5 * p50_10k, f"50K p50 ({p50_50k}) >= 5x 10K p50 ({5 * p50_10k})"
    # 100K should be less than 2.5x the 50K value (sub-linear)
    limit = int(2.5 * p50_50k)
    assert p50_100k < limit, f"100K p50 ({p50_100k}) >= 2.5x 50K p50 ({limit})"
    print("scaling-sublinear-ok: sub-linear scaling verified")
    return 0


def _cmd_run_small_index() -> int:
    """Create a 1K-file fixture, index it, verify throughput."""
    from cleveragents.application.services.repo_indexing_utils import walk_and_index

    tmpdir = tempfile.mkdtemp(prefix="robot-scale-idx-")
    try:
        # Create 1000 files
        for i in range(1000):
            subdir = os.path.join(tmpdir, f"pkg{i % 10}")
            os.makedirs(subdir, exist_ok=True)
            fpath = os.path.join(subdir, f"file_{i:04d}.py")
            with open(fpath, "w") as fh:
                fh.write(f"# module {i}\ndef func_{i}() -> int:\n    return {i}\n")

        t0 = time.perf_counter()
        records = walk_and_index(
            root=Path(tmpdir),
            include_globs=(),
            exclude_globs=("*.pyc",),
            max_file_size=None,
            max_total_size=None,
        )
        elapsed = time.perf_counter() - t0
        file_count = len(records)
        total_tokens = sum(r.token_count for r in records)
        throughput = total_tokens / elapsed if elapsed > 0 else 0.0

        print(f"  files_indexed: {file_count}")
        print(f"  total_tokens: {total_tokens}")
        print(f"  elapsed_s: {elapsed:.3f}")
        print(f"  tokens_per_second: {throughput:.0f}")

        # Minimum threshold: at least 500 files should be indexed
        assert file_count >= 500, f"Only {file_count} files indexed, expected >= 500"
        # Index should complete within the 1K threshold (12 seconds max)
        assert elapsed < 12.0, f"Indexing took {elapsed:.1f}s, exceeds 12s threshold"
        print("scaling-index-ok: small-scale indexing passed")
        return 0
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


# ---------------------------------------------------------------------------
# Main dispatch
# ---------------------------------------------------------------------------


def main() -> int:
    """Entry point called by Robot Framework ``Run Process``."""
    if len(sys.argv) < 2:
        print(
            "Usage: helper_scaling_performance.py "
            "<check-extended-profiles|check-extended-thresholds|"
            "check-context-assembly|check-execution-throughput|"
            "check-sublinear|run-small-index>"
        )
        return 1

    command = sys.argv[1]
    commands: dict[str, Any] = {
        "check-extended-profiles": _cmd_check_extended_profiles,
        "check-extended-thresholds": _cmd_check_extended_thresholds,
        "check-context-assembly": _cmd_check_context_assembly,
        "check-execution-throughput": _cmd_check_execution_throughput,
        "check-sublinear": _cmd_check_sublinear,
        "run-small-index": _cmd_run_small_index,
    }

    handler = commands.get(command)
    if handler is None:
        print(f"Unknown command: {command}")
        return 1

    try:
        result: int = handler()
        return result
    except (AssertionError, KeyError, json.JSONDecodeError) as exc:
        print(f"FAIL: {exc}")
        return 1


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