cleveragents-core/benchmarks/cross_plan_correction_bench.py

"""ASV benchmarks for cross-plan correction cascading performance.

Measures cascade throughput with varying numbers of child plans in
different states: cancel-only, cancel-with-rollback, and mixed states.
"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path
from typing import ClassVar

# Ensure the local *source* tree is importable even when ASV has an
# older build of the package installed.
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

import cleveragents  # noqa: E402

importlib.reload(cleveragents)

from cleveragents.application.services.cross_plan_correction_service import (  # noqa: E402
    CrossPlanCorrectionService,
)
from cleveragents.domain.models.core.correction import (  # noqa: E402
    ChildPlanState,
    CorrectionImpact,
    CorrectionMode,
)

# ---------------------------------------------------------------------------
# Mock implementations (minimal, for benchmarking)
# ---------------------------------------------------------------------------


class _BenchPlanLookup:
    def __init__(self, states: dict[str, ChildPlanState]) -> None:
        self._states = states

    def get_child_plan_state(self, child_plan_id: str) -> ChildPlanState:
        return self._states[child_plan_id]


class _BenchPlanCanceller:
    def cancel_child_plan(self, child_plan_id: str) -> None:
        pass


class _BenchSandboxRollbacker:
    def rollback_child_plan_sandbox(self, child_plan_id: str) -> None:
        pass


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_service(
    child_count: int, state: ChildPlanState
) -> tuple[CrossPlanCorrectionService, list[str]]:
    plan_ids = [f"CP{i}" for i in range(child_count)]
    states = {pid: state for pid in plan_ids}
    lookup = _BenchPlanLookup(states)
    canceller = _BenchPlanCanceller()
    rollbacker = _BenchSandboxRollbacker()
    svc = CrossPlanCorrectionService(
        plan_lookup=lookup,
        plan_canceller=canceller,
        sandbox_rollbacker=rollbacker,
    )
    return svc, plan_ids


# ---------------------------------------------------------------------------
# Benchmark suites
# ---------------------------------------------------------------------------


class CascadeCancelSuite:
    """Benchmark cascade with cancel-only (not-started) child plans."""

    params: ClassVar[list[int]] = [1, 5, 10, 50, 100]
    param_names: ClassVar[list[str]] = ["child_plan_count"]

    def setup(self, child_plan_count: int) -> None:
        self.svc, self.plan_ids = _make_service(
            child_plan_count, ChildPlanState.NOT_STARTED
        )

    def time_execute_cascade_cancel(self, child_plan_count: int) -> None:
        self.svc.execute_cascade("C1", self.plan_ids)

    def time_evaluate_cascade_cancel(self, child_plan_count: int) -> None:
        self.svc.evaluate_cascade("C1", self.plan_ids)


class CascadeRollbackSuite:
    """Benchmark cascade with cancel+rollback (in-progress) child plans."""

    params: ClassVar[list[int]] = [1, 5, 10, 50, 100]
    param_names: ClassVar[list[str]] = ["child_plan_count"]

    def setup(self, child_plan_count: int) -> None:
        self.svc, self.plan_ids = _make_service(
            child_plan_count, ChildPlanState.IN_PROGRESS
        )

    def time_execute_cascade_rollback(self, child_plan_count: int) -> None:
        self.svc.execute_cascade("C1", self.plan_ids)


class CascadeRejectionSuite:
    """Benchmark cascade rejection (applied child plans)."""

    params: ClassVar[list[int]] = [1, 5, 10, 50, 100]
    param_names: ClassVar[list[str]] = ["child_plan_count"]

    def setup(self, child_plan_count: int) -> None:
        self.svc, self.plan_ids = _make_service(
            child_plan_count, ChildPlanState.APPLIED
        )

    def time_execute_cascade_rejection(self, child_plan_count: int) -> None:
        self.svc.execute_cascade("C1", self.plan_ids)


class CorrectionWithCascadeSuite:
    """Benchmark full execute_correction_with_cascade flow."""

    params: ClassVar[list[int]] = [0, 5, 20, 50]
    param_names: ClassVar[list[str]] = ["child_plan_count"]

    def setup(self, child_plan_count: int) -> None:
        plan_ids = [f"CP{i}" for i in range(child_plan_count)]
        states = {pid: ChildPlanState.NOT_STARTED for pid in plan_ids}
        lookup = _BenchPlanLookup(states)
        canceller = _BenchPlanCanceller()
        rollbacker = _BenchSandboxRollbacker()
        self.svc = CrossPlanCorrectionService(
            plan_lookup=lookup,
            plan_canceller=canceller,
            sandbox_rollbacker=rollbacker,
        )
        self.impact = CorrectionImpact(
            affected_decisions=["D1"],
            affected_child_plans=plan_ids,
            risk_level="low",
            rollback_tier="full",
            artifacts_to_archive=["D1.artifact"],
        )

    def time_correction_with_cascade(self, child_plan_count: int) -> None:
        self.svc.execute_correction_with_cascade(
            "C1", self.impact, CorrectionMode.REVERT
        )