cleveragents-core/robot/helper_plan_resume.py

"""Helper script for plan_resume.robot integration tests.

Each sub-command exercises a specific aspect of the plan resume
functionality and prints a success token on completion.
"""

from __future__ import annotations

import sys

from ulid import ULID

from cleveragents.application.services.plan_lifecycle_service import (
    PlanLifecycleService,
)
from cleveragents.application.services.plan_resume_service import (
    PlanResumeService,
)
from cleveragents.config.settings import Settings
from cleveragents.core.exceptions import PlanError, ValidationError
from cleveragents.domain.models.core.plan import (
    ProcessingState,
    ProjectLink,
)
from cleveragents.domain.models.core.resume import (
    ResumeMetadata,
    ResumeSummary,
)


def _make_services() -> tuple[PlanLifecycleService, PlanResumeService]:
    settings = Settings()
    lifecycle = PlanLifecycleService(settings=settings)
    resume = PlanResumeService(lifecycle_service=lifecycle)
    return lifecycle, resume


def _create_action(lifecycle: PlanLifecycleService) -> str:
    name = f"local/resume-robot-{ULID()!s}"[:30].lower()
    lifecycle.create_action(
        name=name,
        description="Robot resume test",
        definition_of_done="Step A\nStep B\nStep C",
        strategy_actor="local/stub-strategy",
        execution_actor="local/stub-execute",
    )
    return name


def _make_errored_plan(
    lifecycle: PlanLifecycleService,
) -> str:
    action = _create_action(lifecycle)
    plan = lifecycle.use_action(
        action_name=action,
        project_links=[ProjectLink(project_name="local/robot-proj")],
    )
    pid = plan.identity.plan_id
    lifecycle.start_strategize(pid)
    p = lifecycle.get_plan(pid)
    p.decision_root_id = str(ULID())
    lifecycle.commit_plan(p)
    lifecycle.complete_strategize(pid)
    lifecycle.execute_plan(pid)
    lifecycle.start_execute(pid)
    lifecycle.fail_execute(pid, "robot test error")
    return pid


def _make_processing_plan(
    lifecycle: PlanLifecycleService,
) -> str:
    action = _create_action(lifecycle)
    plan = lifecycle.use_action(
        action_name=action,
        project_links=[ProjectLink(project_name="local/robot-proj")],
    )
    pid = plan.identity.plan_id
    lifecycle.start_strategize(pid)
    p = lifecycle.get_plan(pid)
    p.decision_root_id = str(ULID())
    lifecycle.commit_plan(p)
    lifecycle.complete_strategize(pid)
    lifecycle.execute_plan(pid)
    lifecycle.start_execute(pid)
    return pid


def _make_applied_plan(
    lifecycle: PlanLifecycleService,
) -> str:
    action = _create_action(lifecycle)
    plan = lifecycle.use_action(
        action_name=action,
        project_links=[ProjectLink(project_name="local/robot-proj")],
    )
    pid = plan.identity.plan_id
    lifecycle.start_strategize(pid)
    p = lifecycle.get_plan(pid)
    p.decision_root_id = str(ULID())
    lifecycle.commit_plan(p)
    lifecycle.complete_strategize(pid)
    lifecycle.execute_plan(pid)
    lifecycle.start_execute(pid)
    lifecycle.complete_execute(pid)
    lifecycle.apply_plan(pid)
    lifecycle.start_apply(pid)
    lifecycle.complete_apply(pid)
    return pid


def test_eligibility_errored() -> None:
    lifecycle, resume = _make_services()
    pid = _make_errored_plan(lifecycle)
    result = resume.validate_eligibility(pid)
    assert result.eligible is True
    print("eligibility-errored-ok")


def test_eligibility_terminal() -> None:
    lifecycle, resume = _make_services()

    # Applied
    pid_applied = _make_applied_plan(lifecycle)
    r1 = resume.validate_eligibility(pid_applied)
    assert r1.eligible is False
    assert r1.ineligible_reason is not None
    assert r1.ineligible_reason.value == "terminal_applied"

    # Cancelled
    pid_cancel = _make_processing_plan(lifecycle)
    lifecycle.cancel_plan(pid_cancel, "test")
    r2 = resume.validate_eligibility(pid_cancel)
    assert r2.eligible is False
    assert r2.ineligible_reason is not None
    assert r2.ineligible_reason.value == "terminal_cancelled"

    print("eligibility-terminal-ok")


def test_record_checkpoints() -> None:
    lifecycle, resume = _make_services()
    pid = _make_processing_plan(lifecycle)

    resume.set_total_steps(pid, 3)
    resume.record_step_checkpoint(pid, 0, "DEC-A", "Step A")
    cp1 = resume.record_step_checkpoint(pid, 1, "DEC-B", "Step B")

    plan = lifecycle.get_plan(pid)
    assert plan.last_completed_step == 1
    assert plan.last_checkpoint_id == cp1.checkpoint_id

    meta = resume.get_resume_metadata(pid)
    assert len(meta.checkpoints) == 2
    assert meta.next_step_index == 2

    print("record-checkpoints-ok")


def test_dry_run_resume() -> None:
    lifecycle, resume = _make_services()
    pid = _make_errored_plan(lifecycle)

    resume.set_total_steps(pid, 5)
    resume.record_step_checkpoint(pid, 2, "DEC-C", "Step C")

    summary = resume.resume_plan(pid, dry_run=True)
    assert summary.next_step_index == 3
    assert summary.total_steps == 5
    assert summary.decision_id == "DEC-C"

    # State should NOT change
    plan = lifecycle.get_plan(pid)
    assert plan.processing_state == ProcessingState.ERRORED

    print("dry-run-resume-ok")


def test_live_resume() -> None:
    lifecycle, resume = _make_services()
    pid = _make_errored_plan(lifecycle)

    resume.set_total_steps(pid, 5)
    resume.record_step_checkpoint(pid, 1, "DEC-B", "Step B")

    summary = resume.resume_plan(pid, dry_run=False)
    assert summary.next_step_index == 2

    plan = lifecycle.get_plan(pid)
    assert plan.processing_state == ProcessingState.PROCESSING
    assert plan.error_message is None

    print("live-resume-ok")


def test_graceful_shutdown() -> None:
    lifecycle, resume = _make_services()
    pid = _make_processing_plan(lifecycle)

    resume.set_total_steps(pid, 3)
    resume.record_step_checkpoint(pid, 0, "DEC-A", "Step A")
    resume.record_shutdown(pid)

    meta = resume.get_resume_metadata(pid)
    assert meta.interrupted is True
    assert meta.interrupted_at is not None

    print("graceful-shutdown-ok")


def test_metadata_properties() -> None:
    # Fresh metadata
    m1 = ResumeMetadata()
    assert m1.has_progress is False
    assert m1.is_complete is False
    assert m1.next_step_index == 0

    # With progress
    m2 = ResumeMetadata(last_completed_step=2, total_steps=5)
    assert m2.has_progress is True
    assert m2.is_complete is False
    assert m2.next_step_index == 3

    # Complete
    m3 = ResumeMetadata(last_completed_step=4, total_steps=5)
    assert m3.has_progress is True
    assert m3.is_complete is True
    assert m3.next_step_index == 5

    print("metadata-properties-ok")


def test_summary_cli_dict() -> None:
    summary = ResumeSummary(
        plan_id="01HGZ6FE0AQDYTR4BXEXAMPLE1",
        phase="execute",
        processing_state="errored",
        last_completed_step=2,
        next_step_index=3,
        total_steps=5,
        decision_id="DEC-001",
        last_checkpoint_id="01HGZ6FE0AQDYTR4BXCHECKPT1",
        sandbox_ref="/tmp/sandbox-test",
    )
    d = summary.as_cli_dict()
    assert d["plan_id"] == "01HGZ6FE0AQDYTR4BXEXAMPLE1"
    assert d["next_step_index"] == 3
    assert d["decision_id"] == "DEC-001"
    assert d["sandbox_ref"] == "/tmp/sandbox-test"

    print("summary-cli-dict-ok")


def test_validation_errors() -> None:
    lifecycle, resume = _make_services()

    # Empty plan_id
    try:
        resume.resume_plan("")
        raise AssertionError("Should have raised")
    except ValidationError:
        pass

    try:
        resume.validate_eligibility("")
        raise AssertionError("Should have raised")
    except ValidationError:
        pass

    try:
        resume.set_total_steps("", 5)
        raise AssertionError("Should have raised")
    except ValidationError:
        pass

    try:
        resume.record_step_checkpoint("", 0, "D", "text")
        raise AssertionError("Should have raised")
    except ValidationError:
        pass

    # Terminal plan
    pid = _make_applied_plan(lifecycle)
    try:
        resume.resume_plan(pid)
        raise AssertionError("Should have raised")
    except PlanError:
        pass

    print("validation-errors-ok")


def test_checkpoint_sandbox() -> None:
    lifecycle, resume = _make_services()
    pid = _make_processing_plan(lifecycle)

    resume.set_total_steps(pid, 3)
    cp = resume.record_step_checkpoint(
        pid, 0, "DEC-A", "Step A", sandbox_ref="/tmp/sandbox-123"
    )
    assert cp.sandbox_ref == "/tmp/sandbox-123"

    summary = resume.resume_plan(pid, dry_run=True)
    assert summary.sandbox_ref == "/tmp/sandbox-123"

    print("checkpoint-sandbox-ok")


TESTS = {
    "eligibility-errored": test_eligibility_errored,
    "eligibility-terminal": test_eligibility_terminal,
    "record-checkpoints": test_record_checkpoints,
    "dry-run-resume": test_dry_run_resume,
    "live-resume": test_live_resume,
    "graceful-shutdown": test_graceful_shutdown,
    "metadata-properties": test_metadata_properties,
    "summary-cli-dict": test_summary_cli_dict,
    "validation-errors": test_validation_errors,
    "checkpoint-sandbox": test_checkpoint_sandbox,
}


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

    test_name = sys.argv[1]
    if test_name not in TESTS:
        print(f"Unknown test: {test_name}", file=sys.stderr)
        print(f"Available: {', '.join(sorted(TESTS))}", file=sys.stderr)
        sys.exit(1)

    TESTS[test_name]()


if __name__ == "__main__":
    main()