cleveragents-core/robot/helper_checkpoint_rollback.py

"""Helper script for checkpoint/rollback Robot Framework smoke tests.

Exercises CheckpointService operations: create, list, rollback,
prune, delete, guards, metadata, and correction-service integration.
"""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

from cleveragents.application.services.checkpoint_service import CheckpointService
from cleveragents.application.services.correction_service import CorrectionService
from cleveragents.core.exceptions import BusinessRuleViolation, ResourceNotFoundError
from cleveragents.domain.models.core.checkpoint import CheckpointRetentionPolicy


def _create_temp_git_workspace() -> str:
    """Create a temporary git workspace with an initial commit."""
    tmpdir = tempfile.mkdtemp(prefix="checkpoint_rollback_test_")
    subprocess.run(
        ["git", "init", "--initial-branch=main", tmpdir],
        check=True,
        capture_output=True,
    )
    subprocess.run(
        ["git", "config", "user.email", "test@example.com"],
        cwd=tmpdir,
        check=True,
        capture_output=True,
    )
    subprocess.run(
        ["git", "config", "user.name", "Test"],
        cwd=tmpdir,
        check=True,
        capture_output=True,
    )
    tracked = os.path.join(tmpdir, "tracked.txt")
    Path(tracked).write_text("initial content\n")
    subprocess.run(
        ["git", "add", "tracked.txt"],
        cwd=tmpdir,
        check=True,
        capture_output=True,
    )
    subprocess.run(
        ["git", "commit", "-m", "Initial commit"],
        cwd=tmpdir,
        check=True,
        capture_output=True,
    )
    return tmpdir


def _get_head_sha(cwd: str) -> str:
    """Return HEAD commit SHA."""
    result = subprocess.run(
        ["git", "rev-parse", "HEAD"],
        cwd=cwd,
        capture_output=True,
        text=True,
        check=True,
    )
    return result.stdout.strip()


_PLAN_ID = "01ARZ3NDEKTSV4RRFFQ69G5FAV"


def _create_checkpoint() -> None:
    svc = CheckpointService()
    cp = svc.create_checkpoint(_PLAN_ID, "commit-abc")
    assert cp.checkpoint_id
    assert cp.plan_id == _PLAN_ID
    print("create-checkpoint-ok")


def _list_checkpoints() -> None:
    svc = CheckpointService()
    svc.create_checkpoint(_PLAN_ID, "commit-1")
    svc.create_checkpoint(_PLAN_ID, "commit-2")
    cps = svc.list_checkpoints(_PLAN_ID)
    assert len(cps) == 2
    print("list-checkpoints-ok")


def _rollback_checkpoint() -> None:
    tmpdir = _create_temp_git_workspace()
    try:
        initial_sha = _get_head_sha(tmpdir)
        svc = CheckpointService()
        svc.register_sandbox(_PLAN_ID, tmpdir)
        cp = svc.create_checkpoint(_PLAN_ID, initial_sha)

        # Modify the tracked file and commit
        Path(os.path.join(tmpdir, "tracked.txt")).write_text("modified\n")
        subprocess.run(
            ["git", "add", "tracked.txt"],
            cwd=tmpdir,
            check=True,
            capture_output=True,
        )
        subprocess.run(
            ["git", "commit", "-m", "Modify tracked file"],
            cwd=tmpdir,
            check=True,
            capture_output=True,
        )

        result = svc.rollback_to_checkpoint(_PLAN_ID, cp.checkpoint_id)
        assert result.restored_files_count > 0
        assert result.from_checkpoint_id == cp.checkpoint_id

        # Verify actual file reversion
        content = Path(os.path.join(tmpdir, "tracked.txt")).read_text()
        assert content == "initial content\n", (
            f"Expected file reverted, got: {content!r}"
        )
        print("rollback-checkpoint-ok")
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def _rollback_applied_guard() -> None:
    svc = CheckpointService()
    svc.register_sandbox(_PLAN_ID, "sandbox-123")
    svc.mark_plan_applied(_PLAN_ID)
    cp = svc.create_checkpoint(_PLAN_ID, "commit-abc")
    try:
        svc.rollback_to_checkpoint(_PLAN_ID, cp.checkpoint_id)
        msg = "Should have raised"
        raise AssertionError(msg)
    except BusinessRuleViolation as e:
        assert "applied" in str(e)
    print("rollback-applied-guard-ok")


def _rollback_sandbox_guard() -> None:
    svc = CheckpointService()
    cp = svc.create_checkpoint(_PLAN_ID, "commit-abc")
    try:
        svc.rollback_to_checkpoint(_PLAN_ID, cp.checkpoint_id)
        msg = "Should have raised"
        raise AssertionError(msg)
    except BusinessRuleViolation as e:
        assert "sandbox" in str(e)
    print("rollback-sandbox-guard-ok")


def _prune_checkpoints() -> None:
    svc = CheckpointService()
    for i in range(5):
        svc.create_checkpoint(_PLAN_ID, f"commit-{i}")
    policy = CheckpointRetentionPolicy(max_checkpoints=3, auto_prune=True)
    pruned = svc.prune_checkpoints(_PLAN_ID, policy)
    assert len(pruned) == 2
    remaining = svc.list_checkpoints(_PLAN_ID)
    assert len(remaining) == 3
    # Verify first and most recent are preserved
    refs = [cp.sandbox_ref for cp in remaining]
    assert "commit-0" in refs, "First checkpoint must be preserved"
    assert "commit-4" in refs, "Most recent checkpoint must be preserved"
    print("prune-checkpoints-ok")


def _delete_checkpoint() -> None:
    svc = CheckpointService()
    cp = svc.create_checkpoint(_PLAN_ID, "commit-abc")
    svc.delete_checkpoint(cp.checkpoint_id)
    try:
        svc.get_checkpoint(cp.checkpoint_id)
        msg = "Should have raised"
        raise AssertionError(msg)
    except ResourceNotFoundError:
        pass
    print("delete-checkpoint-ok")


def _checkpoint_metadata() -> None:
    svc = CheckpointService()
    cp = svc.create_checkpoint(
        _PLAN_ID,
        "commit-abc",
        reason="pre-execution",
        source_tool="lint-check",
        phase="execute",
    )
    stored = svc.get_checkpoint(cp.checkpoint_id)
    assert stored.metadata.reason == "pre-execution"
    assert stored.metadata.source_tool == "lint-check"
    assert stored.metadata.phase == "execute"
    print("checkpoint-metadata-ok")


def _auto_prune_create() -> None:
    svc = CheckpointService()
    for i in range(10):
        svc.create_checkpoint(_PLAN_ID, f"commit-{i}")
    policy = CheckpointRetentionPolicy(max_checkpoints=5, auto_prune=True)
    svc.create_checkpoint(_PLAN_ID, "commit-new", retention_policy=policy)
    remaining = svc.list_checkpoints(_PLAN_ID)
    assert len(remaining) == 5
    print("auto-prune-create-ok")


def _checkpoint_spec_fields() -> None:
    svc = CheckpointService()
    cp = svc.create_checkpoint(
        _PLAN_ID,
        "commit-spec",
        decision_id="01DEC1S10N0000000000000000",
        checkpoint_type="pre_write",
        resource_id="01RES0URCE0000000000000000",
        filesystem_path="snapshots/cp001",
        size_bytes=4096,
    )
    stored = svc.get_checkpoint(cp.checkpoint_id)
    assert stored.decision_id == "01DEC1S10N0000000000000000"
    assert stored.checkpoint_type == "pre_write"
    assert stored.resource_id == "01RES0URCE0000000000000000"
    assert stored.filesystem_path == "snapshots/cp001"
    assert stored.size_bytes == 4096
    print("checkpoint-spec-fields-ok")


def _correction_checkpoint() -> None:
    cp_svc = CheckpointService()
    corr_svc = CorrectionService(checkpoint_service=cp_svc)
    assert corr_svc._checkpoint_service is not None
    print("correction-checkpoint-ok")


_COMMANDS: dict[str, object] = {
    "create-checkpoint": _create_checkpoint,
    "list-checkpoints": _list_checkpoints,
    "rollback-checkpoint": _rollback_checkpoint,
    "rollback-applied-guard": _rollback_applied_guard,
    "rollback-sandbox-guard": _rollback_sandbox_guard,
    "prune-checkpoints": _prune_checkpoints,
    "delete-checkpoint": _delete_checkpoint,
    "checkpoint-metadata": _checkpoint_metadata,
    "auto-prune-create": _auto_prune_create,
    "checkpoint-spec-fields": _checkpoint_spec_fields,
    "correction-checkpoint": _correction_checkpoint,
}

if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        print(f"Usage: {sys.argv[0]} <command>", file=sys.stderr)
        print(f"Commands: {', '.join(sorted(_COMMANDS))}", file=sys.stderr)
        sys.exit(1)

    cmd = _COMMANDS[sys.argv[1]]
    assert callable(cmd)
    cmd()