cleveragents-core/robot/helper_cross_plan_correction.py

"""Helper script for cross-plan correction Robot Framework smoke tests.

Exercises CrossPlanCorrectionService cascade flows, rejection logic,
and atomic rollback without requiring persistence infrastructure.
"""

from __future__ import annotations

import sys

from cleveragents.application.services.cross_plan_correction_service import (
    CrossPlanCorrectionService,
)
from cleveragents.domain.models.core.correction import (
    ChildPlanState,
    CorrectionImpact,
    CorrectionMode,
    CorrectionRejection,
    CorrectionResult,
)

__all__: list[str] = []


# ---------------------------------------------------------------------------
# Mock implementations
# ---------------------------------------------------------------------------


class MockPlanLookup:
    """In-memory child plan state lookup."""

    def __init__(self) -> None:
        self._states: dict[str, ChildPlanState] = {}

    def set_state(self, plan_id: str, state: ChildPlanState) -> None:
        self._states[plan_id] = state

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


class MockPlanCanceller:
    """In-memory child plan canceller."""

    def __init__(self, fail_on: str | None = None) -> None:
        self.cancelled: list[str] = []
        self._fail_on = fail_on

    def cancel_child_plan(self, child_plan_id: str) -> None:
        if self._fail_on and child_plan_id == self._fail_on:
            raise RuntimeError(f"Cancel failed for {child_plan_id}")
        self.cancelled.append(child_plan_id)


class MockSandboxRollbacker:
    """In-memory sandbox rollbacker."""

    def __init__(self) -> None:
        self.rolled_back: list[str] = []

    def rollback_child_plan_sandbox(self, child_plan_id: str) -> None:
        self.rolled_back.append(child_plan_id)


def _make_service(
    states: dict[str, ChildPlanState],
    fail_on: str | None = None,
) -> tuple[
    CrossPlanCorrectionService, MockPlanLookup, MockPlanCanceller, MockSandboxRollbacker
]:
    lookup = MockPlanLookup()
    for pid, s in states.items():
        lookup.set_state(pid, s)
    canceller = MockPlanCanceller(fail_on=fail_on)
    rollbacker = MockSandboxRollbacker()
    svc = CrossPlanCorrectionService(
        plan_lookup=lookup,
        plan_canceller=canceller,
        sandbox_rollbacker=rollbacker,
    )
    return svc, lookup, canceller, rollbacker


# ---------------------------------------------------------------------------
# Command handlers
# ---------------------------------------------------------------------------


def _cancel_not_started() -> None:
    svc, _, _canceller, _rollbacker = _make_service({"CP1": ChildPlanState.NOT_STARTED})
    result = svc.execute_cascade("C1", ["CP1"])
    assert not result.rejected
    assert "CP1" in result.all_cancelled_plan_ids
    assert len(result.all_rolled_back_plan_ids) == 0
    print("cancel-not-started-ok")


def _cancel_rollback_in_progress() -> None:
    svc, _, _canceller, _rollbacker = _make_service({"CP1": ChildPlanState.IN_PROGRESS})
    result = svc.execute_cascade("C1", ["CP1"])
    assert not result.rejected
    assert "CP1" in result.all_cancelled_plan_ids
    assert "CP1" in result.all_rolled_back_plan_ids
    print("cancel-rollback-in-progress-ok")


def _cancel_rollback_completed() -> None:
    svc, _, _, _ = _make_service({"CP1": ChildPlanState.COMPLETED_UNAPPLIED})
    result = svc.execute_cascade("C1", ["CP1"])
    assert not result.rejected
    assert "CP1" in result.all_cancelled_plan_ids
    assert "CP1" in result.all_rolled_back_plan_ids
    print("cancel-rollback-completed-ok")


def _reject_applied() -> None:
    svc, _, _, _ = _make_service({"CP1": ChildPlanState.APPLIED})
    result = svc.execute_cascade("C1", ["CP1"])
    assert result.rejected
    assert result.rejection is not None
    assert "CP1" in result.rejection.affected_applied_child_plan_ids
    print("reject-applied-ok")


def _mixed_states_reject() -> None:
    svc, _, _, _ = _make_service(
        {
            "CP1": ChildPlanState.NOT_STARTED,
            "CP2": ChildPlanState.IN_PROGRESS,
            "CP3": ChildPlanState.APPLIED,
        }
    )
    result = svc.execute_cascade("C1", ["CP1", "CP2", "CP3"])
    assert result.rejected
    assert result.rejection is not None
    assert "CP3" in result.rejection.affected_applied_child_plan_ids
    assert len(result.all_cancelled_plan_ids) == 0
    print("mixed-states-reject-ok")


def _atomic_failure() -> None:
    svc, _, _, _ = _make_service(
        {"CP1": ChildPlanState.NOT_STARTED, "CP2": ChildPlanState.NOT_STARTED},
        fail_on="CP2",
    )
    try:
        svc.execute_cascade("C1", ["CP1", "CP2"])
        raise AssertionError("Expected RuntimeError")
    except RuntimeError as exc:
        assert "CP2" in str(exc)
    print("atomic-failure-ok")


def _cascade_no_children() -> None:
    svc, _, _, _ = _make_service({})
    impact = CorrectionImpact(
        affected_decisions=["D1"],
        affected_child_plans=[],
        risk_level="low",
        rollback_tier="full",
        artifacts_to_archive=["D1.artifact"],
    )
    result = svc.execute_correction_with_cascade("C1", impact, CorrectionMode.REVERT)
    assert isinstance(result, CorrectionResult)
    assert result.status == "applied"
    print("cascade-no-children-ok")


def _cascade_rejection() -> None:
    svc, _, _, _ = _make_service({"CP1": ChildPlanState.APPLIED})
    impact = CorrectionImpact(
        affected_decisions=["D1"],
        affected_child_plans=["CP1"],
        risk_level="low",
        rollback_tier="full",
        artifacts_to_archive=["D1.artifact"],
    )
    result = svc.execute_correction_with_cascade("C1", impact, CorrectionMode.REVERT)
    assert isinstance(result, CorrectionRejection)
    print("cascade-rejection-ok")


# ---------------------------------------------------------------------------
# Dispatcher
# ---------------------------------------------------------------------------

_COMMANDS = {
    "cancel-not-started": _cancel_not_started,
    "cancel-rollback-in-progress": _cancel_rollback_in_progress,
    "cancel-rollback-completed": _cancel_rollback_completed,
    "reject-applied": _reject_applied,
    "mixed-states-reject": _mixed_states_reject,
    "atomic-failure": _atomic_failure,
    "cascade-no-children": _cascade_no_children,
    "cascade-rejection": _cascade_rejection,
}


def main() -> None:
    """Dispatch to the requested command."""
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <command>", file=sys.stderr)
        sys.exit(1)

    command = sys.argv[1]
    handler = _COMMANDS.get(command)
    if handler is None:
        print(f"Unknown command: {command}", file=sys.stderr)
        print(f"Available: {', '.join(sorted(_COMMANDS))}", file=sys.stderr)
        sys.exit(1)

    handler()


if __name__ == "__main__":
    main()