cleveragents-core/robot/helper_revert_re_execution.py

"""Helper script for revert re-execution Robot Framework integration tests.

Exercises the full revert-mode re-execution pipeline: checkpoint restoration,
actor state recovery, phase transition signalling, user intervention decision
creation, and selective subtree recomputation.
"""

from __future__ import annotations

import sys
from unittest.mock import MagicMock

from cleveragents.application.services.checkpoint_service import CheckpointService
from cleveragents.application.services.correction_service import CorrectionService
from cleveragents.domain.models.core.checkpoint import Checkpoint, CheckpointMetadata
from cleveragents.domain.models.core.correction import CorrectionMode
from cleveragents.domain.models.core.decision import (
    ContextSnapshot,
    Decision,
    DecisionType,
)

# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------
_PLAN_ID = "01ARZ3NDEKTSV4RRFFQ69G5FAV"
_DECISION_D1 = "D1"


def _make_decision(
    decision_id: str,
    actor_state_ref: str = "",
) -> Decision:
    """Create a minimal Decision for testing."""
    from ulid import ULID

    return Decision(
        decision_id=str(ULID()),
        plan_id=str(ULID()),
        sequence_number=0,
        decision_type=DecisionType.STRATEGY_CHOICE,
        question="Test question?",
        chosen_option="Option A",
        context_snapshot=ContextSnapshot(actor_state_ref=actor_state_ref),
    )


def _mock_checkpoint_service(
    checkpoints: list[Checkpoint] | None = None,
) -> CheckpointService:
    """Create a mock CheckpointService."""
    mock = MagicMock(spec=CheckpointService)
    mock.list_checkpoints.return_value = checkpoints or []
    mock.rollback_to_checkpoint.return_value = MagicMock(
        restored_files_count=3,
        changed_paths=["a.py", "b.py", "c.py"],
        from_checkpoint_id="cp1",
    )
    return mock


def _make_checkpoint(
    plan_id: str,
    decision_id: str,
    sandbox_ref: str,
) -> Checkpoint:
    """Create a Checkpoint aligned to a decision."""
    from ulid import ULID

    return Checkpoint(
        checkpoint_id=str(ULID()),
        plan_id=plan_id,
        sandbox_ref=sandbox_ref,
        decision_id=decision_id,
        checkpoint_type="pre_write",
        metadata=CheckpointMetadata(reason="decision-aligned"),
    )


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


def _revert_with_checkpoint() -> None:
    cp = _make_checkpoint(_PLAN_ID, _DECISION_D1, "abc123")
    mock_cp_svc = _mock_checkpoint_service([cp])
    svc = CorrectionService(checkpoint_service=mock_cp_svc)
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id)
    assert result.checkpoint_restored is True, (
        f"Expected checkpoint_restored=True, got {result.checkpoint_restored}"
    )
    assert result.status == "applied"
    mock_cp_svc.rollback_to_checkpoint.assert_called_once()
    print("revert-with-checkpoint-ok")


def _revert_no_checkpoint() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id)
    assert result.checkpoint_restored is False
    assert result.status == "applied"
    print("revert-no-checkpoint-ok")


def _revert_actor_state() -> None:
    decision = _make_decision(_DECISION_D1, actor_state_ref="lg-state-42")
    decisions = {_DECISION_D1: decision}
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id, decisions=decisions)
    assert result.actor_state_ref == "lg-state-42", (
        f"Expected actor_state_ref='lg-state-42', got {result.actor_state_ref!r}"
    )
    print("revert-actor-state-ok")


def _revert_phase_transition() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id)
    assert result.phase_transition_target == "strategize", (
        f"Expected phase_transition_target='strategize', "
        f"got {result.phase_transition_target!r}"
    )
    print("revert-phase-transition-ok")


def _revert_intervention_id() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id)
    assert result.user_intervention_decision_id, (
        "Expected non-empty user_intervention_decision_id"
    )
    print("revert-intervention-id-ok")


def _full_revert_flow() -> None:
    cp = _make_checkpoint(_PLAN_ID, _DECISION_D1, "commit-sha-1")
    mock_cp_svc = _mock_checkpoint_service([cp])
    decision = _make_decision(_DECISION_D1, actor_state_ref="lg-state-99")
    decisions = {_DECISION_D1: decision}
    tree = {_DECISION_D1: ["D2", "D3"]}

    svc = CorrectionService(checkpoint_service=mock_cp_svc)
    req = svc.request_correction(
        _PLAN_ID,
        _DECISION_D1,
        CorrectionMode.REVERT,
        guidance="Prioritize payments",
    )
    result = svc.execute_revert(
        req.correction_id,
        decision_tree=tree,
        decisions=decisions,
    )

    assert result.status == "applied"
    assert result.checkpoint_restored is True
    assert result.actor_state_ref == "lg-state-99"
    assert result.phase_transition_target == "strategize"
    assert result.user_intervention_decision_id
    assert "D1" in result.reverted_decisions
    assert "D2" in result.reverted_decisions
    assert "D3" in result.reverted_decisions
    print("full-revert-flow-ok")


def _revert_selective_subtree() -> None:
    tree = {"ROOT": ["D1", "D2"], "D1": ["D3"]}
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, "D1", CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id, decision_tree=tree)
    assert "D1" in result.reverted_decisions
    assert "D3" in result.reverted_decisions
    assert "ROOT" not in result.reverted_decisions
    assert "D2" not in result.reverted_decisions
    print("revert-selective-subtree-ok")


# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

COMMANDS: dict[str, object] = {
    "revert-with-checkpoint": _revert_with_checkpoint,
    "revert-no-checkpoint": _revert_no_checkpoint,
    "revert-actor-state": _revert_actor_state,
    "revert-phase-transition": _revert_phase_transition,
    "revert-intervention-id": _revert_intervention_id,
    "full-revert-flow": _full_revert_flow,
    "revert-selective-subtree": _revert_selective_subtree,
}


def main() -> None:
    """Dispatch command from sys.argv."""
    if len(sys.argv) < 2:
        raise SystemExit("Expected command argument")
    command = sys.argv[1]
    if command not in COMMANDS:
        raise SystemExit(f"Unknown command: {command}")
    func = COMMANDS[command]
    if callable(func):
        func()


if __name__ == "__main__":
    main()