cleveragents-core/robot/helper_correction_attempt_persistence.py

"""Helper script for Robot Framework correction attempt persistence tests.

Usage:
    python robot/helper_correction_attempt_persistence.py <subcommand>

Subcommands:
    create-retrieve   Create + retrieve a correction attempt
    list-by-plan      Create multiple, list by plan ID
    update-state      Update state transitions
    delete            Delete a correction attempt
    domain-roundtrip  Verify all spec DDL columns round-trip
"""

from __future__ import annotations

import sys
from datetime import UTC, datetime, timedelta
from pathlib import Path

# Ensure src is importable when run from workspace root
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from collections.abc import Callable
from typing import Any

from sqlalchemy import create_engine, event
from sqlalchemy.orm import Session, sessionmaker

from cleveragents.domain.models.core.action import (
    Action,
    ActionState,
)
from cleveragents.domain.models.core.correction import (
    CorrectionAttemptRecord,
    CorrectionAttemptState,
    CorrectionMode,
)
from cleveragents.domain.models.core.decision import (
    ContextSnapshot,
    Decision,
    DecisionType,
)
from cleveragents.domain.models.core.plan import (
    NamespacedName,
    Plan,
    PlanIdentity,
    PlanPhase,
    PlanTimestamps,
    ProcessingState,
)
from cleveragents.infrastructure.database.models import Base
from cleveragents.infrastructure.database.repositories import (
    ActionRepository,
    CorrectionAttemptNotFoundError,
    CorrectionAttemptRepository,
    DecisionRepository,
    LifecyclePlanRepository,
)

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

_PLAN_ID = "01HV000000000000000000RC01"
_DECISION_ID = "01HV000000000000000000RD01"


def _setup() -> tuple[Session, Callable[[], Session]]:
    """Create an in-memory SQLite DB with all tables and FK enforcement."""
    engine = create_engine("sqlite:///:memory:", echo=False)

    @event.listens_for(engine, "connect")
    def _enable_fk(dbapi_conn: Any, _rec: Any) -> None:
        cursor = dbapi_conn.cursor()
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()

    Base.metadata.create_all(engine)
    sm = sessionmaker(bind=engine)
    session = sm()
    factory: Callable[[], Session] = lambda: session  # noqa: E731
    return session, factory


def _create_prerequisites(session: Session, factory: Callable[[], Session]) -> None:
    """Create action + plan + decision needed to satisfy FK constraints."""
    now = datetime.now(UTC)

    action_repo = ActionRepository(session_factory=factory)
    action = Action(
        namespaced_name=NamespacedName.parse("local/robot-correction-action"),
        description="Robot correction test action",
        definition_of_done="Correction is applied",
        strategy_actor="strategy-actor",
        execution_actor="execution-actor",
        state=ActionState.AVAILABLE,
    )
    action_repo.create(action)
    session.commit()

    plan_repo = LifecyclePlanRepository(session_factory=factory)
    plan = Plan(
        identity=PlanIdentity(plan_id=_PLAN_ID, attempt=1),
        namespaced_name=NamespacedName.parse("local/robot-test-plan"),
        action_name="local/robot-correction-action",
        description="Robot test plan for correction persistence",
        phase=PlanPhase.EXECUTE,
        processing_state=ProcessingState.PROCESSING,
        timestamps=PlanTimestamps(created_at=now, updated_at=now),
    )
    plan_repo.create(plan)
    session.commit()

    dec_repo = DecisionRepository(session_factory=factory)
    decision = Decision(
        decision_id=_DECISION_ID,
        plan_id=_PLAN_ID,
        sequence_number=0,
        decision_type=DecisionType.PROMPT_DEFINITION,
        question="What approach?",
        chosen_option="Build REST API",
        context_snapshot=ContextSnapshot(
            hot_context_hash="sha256:robot",
            hot_context_ref="ref:robot",
            relevant_resources=[],
            actor_state_ref="",
        ),
    )
    dec_repo.create(decision)
    session.commit()


# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------


def cmd_create_retrieve() -> None:
    """Create a correction attempt and retrieve it."""
    session, factory = _setup()
    _create_prerequisites(session, factory)
    repo = CorrectionAttemptRepository(session_factory=factory)

    record = CorrectionAttemptRecord(
        plan_id=_PLAN_ID,
        original_decision_id=_DECISION_ID,
        mode=CorrectionMode.REVERT,
        guidance="Fix the broken implementation",
    )
    created = repo.create(record)
    session.commit()

    retrieved = repo.get(created.correction_attempt_id)
    assert retrieved.correction_attempt_id == created.correction_attempt_id
    assert retrieved.mode == CorrectionMode.REVERT
    assert retrieved.state == CorrectionAttemptState.PENDING
    assert retrieved.plan_id == _PLAN_ID
    assert retrieved.original_decision_id == _DECISION_ID
    print("create-retrieve-ok")


def cmd_list_by_plan() -> None:
    """Create multiple correction attempts and list by plan."""
    session, factory = _setup()
    _create_prerequisites(session, factory)
    repo = CorrectionAttemptRepository(session_factory=factory)

    base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC)
    for i in range(3):
        record = CorrectionAttemptRecord(
            plan_id=_PLAN_ID,
            original_decision_id=_DECISION_ID,
            mode=CorrectionMode.REVERT,
            guidance=f"Guidance {i}",
            created_at=base_time + timedelta(seconds=i),
        )
        repo.create(record)
        session.commit()

    results = repo.list_by_plan(_PLAN_ID)
    assert len(results) == 3
    # Check ordering
    for i in range(len(results) - 1):
        assert results[i].created_at <= results[i + 1].created_at
    print("list-by-plan-ok")


def cmd_update_state() -> None:
    """Update state transitions."""
    session, factory = _setup()
    _create_prerequisites(session, factory)
    repo = CorrectionAttemptRepository(session_factory=factory)

    record = CorrectionAttemptRecord(
        plan_id=_PLAN_ID,
        original_decision_id=_DECISION_ID,
        mode=CorrectionMode.REVERT,
        guidance="Fix it",
    )
    created = repo.create(record)
    session.commit()

    # Update to executing
    updated = repo.update_state(
        created.correction_attempt_id,
        state=CorrectionAttemptState.EXECUTING,
    )
    session.commit()
    assert updated.state == CorrectionAttemptState.EXECUTING

    # Update to complete with timestamp
    updated = repo.update_state(
        created.correction_attempt_id,
        state=CorrectionAttemptState.COMPLETE,
        completed_at=datetime.now(UTC),
    )
    session.commit()
    assert updated.state == CorrectionAttemptState.COMPLETE
    assert updated.completed_at is not None
    print("update-state-ok")


def cmd_delete() -> None:
    """Delete a correction attempt."""
    session, factory = _setup()
    _create_prerequisites(session, factory)
    repo = CorrectionAttemptRepository(session_factory=factory)

    record = CorrectionAttemptRecord(
        plan_id=_PLAN_ID,
        original_decision_id=_DECISION_ID,
        mode=CorrectionMode.REVERT,
        guidance="Delete me",
    )
    created = repo.create(record)
    session.commit()

    result = repo.delete(created.correction_attempt_id)
    session.commit()
    assert result is True

    # Verify deleted
    try:
        repo.get(created.correction_attempt_id)
        msg = "Should have raised CorrectionAttemptNotFoundError"
        raise AssertionError(msg)
    except CorrectionAttemptNotFoundError:
        pass

    # Delete non-existent
    result = repo.delete("01HV000000000000000NONEXIST")
    assert result is False
    print("delete-ok")


def cmd_domain_roundtrip() -> None:
    """Verify all spec DDL columns survive domain-model round-trip."""
    session, factory = _setup()
    _create_prerequisites(session, factory)
    repo = CorrectionAttemptRepository(session_factory=factory)

    now = datetime.now(UTC)
    record = CorrectionAttemptRecord(
        plan_id=_PLAN_ID,
        original_decision_id=_DECISION_ID,
        mode=CorrectionMode.APPEND,
        guidance="Detailed guidance text for correction",
        archived_artifacts_path="/tmp/archived/artifacts",
        state=CorrectionAttemptState.PENDING,
        created_at=now,
    )
    created = repo.create(record)
    session.commit()

    retrieved = repo.get(created.correction_attempt_id)

    # Verify all columns
    assert retrieved.correction_attempt_id == created.correction_attempt_id
    assert retrieved.plan_id == _PLAN_ID
    assert retrieved.original_decision_id == _DECISION_ID
    assert retrieved.new_decision_id is None
    assert retrieved.mode == CorrectionMode.APPEND
    assert retrieved.guidance == "Detailed guidance text for correction"
    assert retrieved.archived_artifacts_path == "/tmp/archived/artifacts"
    assert retrieved.state == CorrectionAttemptState.PENDING
    assert retrieved.created_at is not None
    # Verify created_at value survives the round-trip within 10ms
    # (from_domain truncates to millisecond precision)
    delta = abs((retrieved.created_at - now).total_seconds())
    assert delta < 0.01, f"created_at drift {delta}s exceeds 10ms tolerance"
    assert retrieved.completed_at is None
    print("domain-roundtrip-ok")


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

_COMMANDS = {
    "create-retrieve": cmd_create_retrieve,
    "list-by-plan": cmd_list_by_plan,
    "update-state": cmd_update_state,
    "delete": cmd_delete,
    "domain-roundtrip": cmd_domain_roundtrip,
}


def main() -> None:
    """Dispatch to the requested subcommand."""
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        print(f"Usage: {sys.argv[0]} <{'|'.join(_COMMANDS)}>", file=sys.stderr)
        sys.exit(1)
    _COMMANDS[sys.argv[1]]()


if __name__ == "__main__":
    main()