"""Helper script for Robot Framework correction attempt persistence tests. Usage: python robot/helper_correction_attempt_persistence.py 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()