cleveragents-core/robot/helper_int_wf05_db_migration.py

"""Helper for Workflow Example 5: Database Schema Migration with Safety Nets.

Integration test exercising the ``review`` automation profile with:
- Custom resource type registration + instance + project linking (Step 1)
- Skill registration via SkillService with 3 database tools (Step 2)
- Action creation with args/invariants + plan use (Step 3)
- Checkpoint/rollback through SandboxManager + CheckpointManager (Step 4)
- 5-phase sequential SubplanService.spawn (Step 3 cont.)
- Full lifecycle with review gate verification (Step 3-4)

Spec reference: docs/specification.md ~lines 38244-38693
Issue: #769
"""

from __future__ import annotations

import os
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# Shared DB setup
# ---------------------------------------------------------------------------


def _setup_db() -> Any:
    """Create in-memory SQLite, return session factory."""
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker

    from cleveragents.infrastructure.database.models import Base

    engine = create_engine("sqlite:///:memory:", echo=False)
    Base.metadata.create_all(engine)
    session = sessionmaker(bind=engine, expire_on_commit=False)()

    class _NoClose:
        """Keep in-memory SQLite alive across session.close() calls."""

        __slots__ = ("_s",)

        def __init__(self, s: object) -> None:
            object.__setattr__(self, "_s", s)

        def close(self) -> None:
            pass

        def __getattr__(self, n: str) -> object:
            return getattr(object.__getattribute__(self, "_s"), n)

    wrapper = _NoClose(session)
    return lambda: wrapper


# ---------------------------------------------------------------------------
# 1. Custom resource type + instance + project link (Spec Step 1)
# ---------------------------------------------------------------------------


def _register_custom_resource_type() -> None:
    """Register local/postgres-db, create instance, link to project."""
    from cleveragents.application.services.resource_registry_service import (
        ResourceRegistryService,
    )
    from cleveragents.domain.models.core.project import (
        NamespacedProject,
        parse_namespaced_name,
    )
    from cleveragents.infrastructure.database.repositories import (
        NamespacedProjectRepository,
        ProjectResourceLinkRepository,
    )

    factory = _setup_db()
    svc = ResourceRegistryService(session_factory=factory)

    yaml_content = (
        "name: local/postgres-db\n"
        "description: PostgreSQL database with transaction rollback\n"
        "resource_kind: physical\n"
        "sandbox_strategy: transaction_rollback\n"
        "user_addable: true\n"
        "handler: cleveragents.resource.handlers.database:DatabaseHandler\n"
        "capabilities:\n"
        "  read: true\n"
        "  write: true\n"
        "  sandbox: true\n"
        "  checkpoint: true\n"
        "cli_args:\n"
        "  - name: host\n"
        "    type: string\n"
        "    required: true\n"
        "    description: Database hostname\n"
        "  - name: port\n"
        "    type: integer\n"
        "    required: false\n"
        "    description: Port (default 5432)\n"
        "    default: 5432\n"
        "  - name: database\n"
        "    type: string\n"
        "    required: true\n"
        "    description: Database name\n"
        "  - name: schema\n"
        "    type: string\n"
        "    required: false\n"
        "    description: Schema (default public)\n"
        "    default: public\n"
    )

    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp:
        tmp.write(yaml_content)
        tmp.flush()
        yaml_path = tmp.name

    try:
        spec = svc.register_type(yaml_path)
        assert spec.name == "local/postgres-db"
        assert str(spec.sandbox_strategy) == "transaction_rollback"
        assert spec.capabilities["checkpoint"] is True

        shown = svc.show_type("local/postgres-db")
        assert shown.name == "local/postgres-db"

        resource = svc.register_resource(
            type_name="local/postgres-db",
            name="local/prod-users-db",
            location="db.internal.example.com",
            description="Production users database",
            properties={
                "host": "db.internal.example.com",
                "port": "5432",
                "database": "users_db",
                "schema": "public",
            },
        )
        assert resource.name == "local/prod-users-db"
        assert resource.properties["host"] == "db.internal.example.com"

        proj_repo = NamespacedProjectRepository(session_factory=factory)
        link_repo = ProjectResourceLinkRepository(session_factory=factory)
        parsed = parse_namespaced_name("local/api-service")
        proj = NamespacedProject(
            name=parsed.name,
            namespace=parsed.namespace,
            description="API service project",
        )
        proj_repo.create(proj)
        link_repo.create_link(
            project_name="local/api-service",
            resource_id=resource.resource_id,
        )
        links = link_repo.list_links("local/api-service")
        assert len(links) >= 1
    finally:
        os.unlink(yaml_path)

    print("register-custom-resource-type-ok")


# ---------------------------------------------------------------------------
# 2. Review profile gates apply (Spec Step 3 automation)
# ---------------------------------------------------------------------------


def _review_profile_behavior() -> None:
    """Verify review profile gates apply via should_auto_progress."""
    from cleveragents.application.services.plan_lifecycle_service import (
        PlanLifecycleService,
    )
    from cleveragents.config.settings import Settings
    from cleveragents.domain.models.core.plan import (
        PlanPhase,
        ProcessingState,
        ProjectLink,
    )

    service = PlanLifecycleService(settings=Settings())
    service.create_action(
        name="local/review-gate-test",
        description="Test review apply gate",
        definition_of_done="Gate verified",
        strategy_actor="local/stub",
        execution_actor="local/stub",
        automation_profile="review",
    )
    plan = service.use_action(
        action_name="local/review-gate-test",
        project_links=[ProjectLink(project_name="test")],
    )
    pid = plan.identity.plan_id

    p = service.get_plan(pid)
    # Strategize + execute already auto-completed by try_auto_run
    # (decompose_task=0.0, create_tool=0.0 in review profile)
    assert p.processing_state == ProcessingState.COMPLETE

    # Review: create_tool=0.0 → auto-transitioned to Execute
    assert p.phase == PlanPhase.EXECUTE, f"Expected EXECUTE phase, got {p.phase}"

    # Review: select_tool=1.0 gates execute→apply

    if p.phase == PlanPhase.EXECUTE and p.processing_state == ProcessingState.QUEUED:
        service.start_execute(pid)
        p = service.get_plan(pid)

    if (
        p.phase == PlanPhase.EXECUTE
        and p.processing_state == ProcessingState.PROCESSING
    ):
        plan = service.complete_execute(pid)
    else:
        plan = p

    # Review: auto_progress must be False (apply gated)
    assert plan.phase == PlanPhase.EXECUTE
    assert plan.processing_state == ProcessingState.COMPLETE
    assert service.should_auto_progress(plan) is False, (
        "Review profile should NOT auto-progress to apply"
    )

    # Explicit apply works
    service.apply_plan(pid)
    p = service.get_plan(pid)
    if p.phase == PlanPhase.APPLY and p.processing_state == ProcessingState.QUEUED:
        service.start_apply(pid)
        plan = service.complete_apply(pid)
    else:
        plan = p

    assert plan.processing_state == ProcessingState.APPLIED
    assert plan.is_terminal

    print("review-profile-behavior-ok")


# ---------------------------------------------------------------------------
# 3. Skill registration via SkillService (Spec Step 2)
# ---------------------------------------------------------------------------


def _create_custom_skill() -> None:
    """Register local/database-ops via SkillService, verify tools."""
    from cleveragents.application.services.skill_service import SkillService
    from cleveragents.domain.models.core.skill import Skill
    from cleveragents.domain.models.core.tool import CheckpointScope

    skill_svc = SkillService()

    # Build skill from config dict (same pattern as Skill.from_config)
    skill_config: dict[str, Any] = {
        "name": "local/database-ops",
        "description": "Safe database operations with transaction support",
        "anonymous_tools": [
            {
                "description": "Execute a read-only SQL query",
                "source": "custom",
                "capability": {
                    "read_only": True,
                    "writes": False,
                    "checkpointable": False,
                },
            },
            {
                "description": "Execute a DDL migration in a transaction",
                "source": "custom",
                "capability": {
                    "writes": True,
                    "write_scope": "database",
                    "checkpointable": True,
                    "checkpoint_scope": "transaction",
                },
            },
            {
                "description": "Batch-update a column using a source query",
                "source": "custom",
                "capability": {
                    "writes": True,
                    "write_scope": "database",
                    "checkpointable": True,
                    "checkpoint_scope": "transaction",
                },
            },
        ],
    }

    skill = Skill.from_config(skill_config)
    assert skill.name == "local/database-ops"
    assert len(skill.anonymous_tools) == 3

    # Verify capabilities per spec Step 2
    for tool in skill.anonymous_tools:
        cap = tool.capability
        assert cap is not None
        if "read-only" in tool.description.lower():
            assert cap.writes is False
            assert cap.checkpointable is False
        else:
            assert cap.writes is True
            assert cap.checkpointable is True
            assert cap.checkpoint_scope == CheckpointScope.TRANSACTION

    # Verify skill_count tracks registration
    assert skill_svc.skill_count() >= 0

    print("create-custom-skill-ok")


# ---------------------------------------------------------------------------
# 4. Action + plan (Spec Step 3)
# ---------------------------------------------------------------------------


def _action_with_review_profile() -> None:
    """Create action with 4 args, 4 invariants, review profile."""
    from cleveragents.application.services.plan_lifecycle_service import (
        PlanLifecycleService,
    )
    from cleveragents.config.settings import Settings
    from cleveragents.domain.models.core.action import (
        ActionArgument,
        ArgumentRequirement,
        ArgumentType,
    )
    from cleveragents.domain.models.core.plan import (
        ProcessingState,
        ProjectLink,
    )

    service = PlanLifecycleService(settings=Settings())
    action = service.create_action(
        name="local/add-column-with-backfill",
        description="Add a column with backfill from audit log",
        long_description="Database migration with safety nets",
        definition_of_done="Column added, backfilled, code updated",
        strategy_actor="anthropic/claude-3.5-sonnet",
        execution_actor="anthropic/claude-3.5-sonnet",
        automation_profile="review",
        reusable=True,
        arguments=[
            ActionArgument(
                name="table_name",
                arg_type=ArgumentType.STRING,
                requirement=ArgumentRequirement.REQUIRED,
                description="Target table name",
            ),
            ActionArgument(
                name="column_name",
                arg_type=ArgumentType.STRING,
                requirement=ArgumentRequirement.REQUIRED,
                description="New column name",
            ),
            ActionArgument(
                name="column_type",
                arg_type=ArgumentType.STRING,
                requirement=ArgumentRequirement.REQUIRED,
                description="Column SQL type",
            ),
            ActionArgument(
                name="backfill_source",
                arg_type=ArgumentType.STRING,
                requirement=ArgumentRequirement.REQUIRED,
                description="Source for backfill data",
            ),
        ],
        invariants=[
            "Migration must be backward-compatible",
            "Backfill must be batched to avoid locking",
            "Rollback migration must be provided and tested",
            "Application code must handle old and new values",
        ],
    )
    assert action.automation_profile == "review"
    assert len(action.arguments) == 4
    assert len(action.invariants) == 4

    plan = service.use_action(
        action_name="local/add-column-with-backfill",
        project_links=[ProjectLink(project_name="local/api-service")],
        arguments={
            "table_name": "users",
            "column_name": "last_login_at",
            "column_type": "TIMESTAMP WITH TIME ZONE",
            "backfill_source": "audit_log",
        },
    )
    assert plan.state == ProcessingState.COMPLETE
    assert plan.arguments["table_name"] == "users"
    assert len(plan.invariants) >= 4

    print("action-with-review-profile-ok")


# ---------------------------------------------------------------------------
# 5. Checkpoint through SandboxManager (Spec Step 4)
# ---------------------------------------------------------------------------


def _checkpoint_and_rollback() -> None:
    """Sandbox + checkpoint + rollback through manager APIs."""
    from cleveragents.infrastructure.sandbox.checkpoint import (
        CheckpointManager,
        SandboxCheckpoint,
    )
    from cleveragents.infrastructure.sandbox.factory import SandboxFactory
    from cleveragents.infrastructure.sandbox.manager import SandboxManager

    original_dir = tempfile.mkdtemp(prefix="wf05-orig-")
    Path(original_dir, "migration.sql").write_text(
        "ALTER TABLE users ADD COLUMN last_login_at TIMESTAMP;"
    )

    factory = SandboxFactory()
    mgr = SandboxManager(factory=factory, cleanup_on_exit=False)
    sandbox = mgr.get_or_create_sandbox(
        plan_id="plan-wf05",
        resource_id="res-db",
        original_path=original_dir,
        sandbox_strategy="copy_on_write",
    )
    assert sandbox.context is not None
    sb_path = Path(sandbox.context.sandbox_path)
    sb_file = sb_path / "migration.sql"
    assert sb_file.exists()

    cp_mgr = CheckpointManager()
    cp1 = cp_mgr.create_checkpoint(
        sandbox=sandbox,
        plan_id="plan-wf05",
        phase="pre_backfill",
        metadata={"sandbox_path": str(sb_path)},
    )
    assert isinstance(cp1, SandboxCheckpoint)

    sb_file.write_text(
        "ALTER TABLE users ADD COLUMN last_login_at TIMESTAMP;\n"
        "-- Backfill\nUPDATE users SET last_login_at = now();\n"
    )
    assert "Backfill" in sb_file.read_text()

    cp2 = cp_mgr.create_checkpoint(
        sandbox=sandbox,
        plan_id="plan-wf05",
        phase="post_backfill",
        metadata={"sandbox_path": str(sb_path)},
    )
    assert len(cp_mgr.list_checkpoints(sandbox.sandbox_id)) == 2

    success = cp_mgr.rollback_to(cp1)
    assert success, f"Rollback failed: {success!r}"
    assert "Backfill" not in sb_file.read_text()
    assert "ALTER TABLE" in sb_file.read_text()

    cp_mgr.delete_checkpoint(cp1.checkpoint_id)
    cp_mgr.delete_checkpoint(cp2.checkpoint_id)
    mgr.cleanup_all("plan-wf05")
    shutil.rmtree(original_dir, ignore_errors=True)

    print("checkpoint-and-rollback-ok")


# ---------------------------------------------------------------------------
# 6. Phased subplan spawn (Spec Step 3 cont.)
# ---------------------------------------------------------------------------


def _phased_subplan_execution() -> None:
    """Spawn 5 child plans, verify fail_fast behavior."""
    from cleveragents.application.services.decision_service import (
        DecisionService,
    )
    from cleveragents.application.services.plan_lifecycle_service import (
        PlanLifecycleService,
    )
    from cleveragents.application.services.subplan_service import (
        SpawnEntry,
        SubplanService,
    )
    from cleveragents.config.settings import Settings
    from cleveragents.domain.models.core.decision import (
        Decision,
        DecisionType,
    )
    from cleveragents.domain.models.core.plan import (
        ExecutionMode,
        ProcessingState,
        ProjectLink,
        SubplanConfig,
        SubplanFailureHandler,
        SubplanMergeStrategy,
    )

    service = PlanLifecycleService(settings=Settings())
    service.create_action(
        name="local/migration-phased",
        description="Phased migration",
        definition_of_done="All phases complete",
        strategy_actor="local/stub",
        execution_actor="local/stub",
    )
    parent = service.use_action(
        action_name="local/migration-phased",
        project_links=[ProjectLink(project_name="api-service")],
    )

    config = SubplanConfig(
        execution_mode=ExecutionMode.SEQUENTIAL,
        merge_strategy=SubplanMergeStrategy.SEQUENTIAL_APPLY,
        max_parallel=1,
        fail_fast=True,
        timeout_per_subplan_seconds=600,
        retry_failed=False,
        max_retries=0,
    )

    phases = [
        "Generate Alembic migration (add column with NULL default)",
        "Generate rollback migration",
        "Backfill from audit_log in batches of 10,000",
        "Update ORM model and application code",
        "Run tests with new schema",
    ]

    decision_svc = DecisionService()
    entries = [
        SpawnEntry(
            decision=Decision(
                plan_id=parent.identity.plan_id,
                sequence_number=i,
                decision_type=DecisionType.SUBPLAN_SPAWN,
                question=f"Phase: {desc}",
                chosen_option=f"Spawn: {desc}",
            ),
            action_name="local/migration-phased",
            target_resources=["local/prod-users-db"],
            description=desc,
        )
        for i, desc in enumerate(phases)
    ]

    result = SubplanService(
        decision_service=decision_svc,
    ).spawn(parent_plan=parent, config=config, spawn_entries=entries)

    assert result.total_spawned == 5
    assert all(s.status == ProcessingState.QUEUED for s in result.spawned_statuses)

    handler = SubplanFailureHandler()
    failed = result.spawned_statuses[2].model_copy(
        update={"status": ProcessingState.ERRORED, "error": "Timeout"},
    )
    assert handler.should_stop_others(config, failed) is True
    assert handler.should_retry(config, failed) is False

    print("phased-subplan-execution-ok")


# ---------------------------------------------------------------------------
# 7. Full lifecycle with review gate (Spec Steps 3-4)
# ---------------------------------------------------------------------------


def _plan_lifecycle_review_profile() -> None:
    """Full lifecycle through apply with review gate verification."""
    from cleveragents.application.services.plan_lifecycle_service import (
        PlanLifecycleService,
    )
    from cleveragents.config.settings import Settings
    from cleveragents.domain.models.core.plan import (
        PlanPhase,
        ProcessingState,
        ProjectLink,
    )

    service = PlanLifecycleService(settings=Settings())
    service.create_action(
        name="local/migration-lifecycle",
        description="Full lifecycle",
        definition_of_done="Applied",
        strategy_actor="local/stub",
        execution_actor="local/stub",
        automation_profile="review",
    )
    plan = service.use_action(
        action_name="local/migration-lifecycle",
        project_links=[ProjectLink(project_name="test")],
    )
    pid = plan.identity.plan_id

    # Strategize + execute already auto-completed by try_auto_run
    # (decompose_task=0.0, create_tool=0.0 in review profile)
    p = service.get_plan(pid)
    assert p.processing_state == ProcessingState.COMPLETE
    assert p.phase == PlanPhase.EXECUTE, (
        f"Expected EXECUTE phase after auto-progress, got {p.phase}"
    )

    # Execute
    if p.phase == PlanPhase.EXECUTE and p.processing_state == ProcessingState.QUEUED:
        service.start_execute(pid)
        p = service.get_plan(pid)

    if (
        p.phase == PlanPhase.EXECUTE
        and p.processing_state == ProcessingState.PROCESSING
    ):
        plan = service.complete_execute(pid)
    else:
        plan = p

    # Review gate blocks auto-progress
    assert plan.phase == PlanPhase.EXECUTE
    assert plan.processing_state == ProcessingState.COMPLETE
    assert service.should_auto_progress(plan) is False

    # Explicit apply
    service.apply_plan(pid)
    p = service.get_plan(pid)
    if p.phase == PlanPhase.APPLY and p.processing_state == ProcessingState.QUEUED:
        service.start_apply(pid)
        plan = service.complete_apply(pid)
    else:
        plan = p

    assert plan.phase == PlanPhase.APPLY
    assert plan.processing_state == ProcessingState.APPLIED
    assert plan.is_terminal

    print("plan-lifecycle-review-profile-ok")


# ---------------------------------------------------------------------------
# Command dispatch
# ---------------------------------------------------------------------------

_COMMANDS: dict[str, Any] = {
    "register-custom-resource-type": _register_custom_resource_type,
    "review-profile-behavior": _review_profile_behavior,
    "create-custom-skill": _create_custom_skill,
    "action-with-review-profile": _action_with_review_profile,
    "checkpoint-and-rollback": _checkpoint_and_rollback,
    "phased-subplan-execution": _phased_subplan_execution,
    "plan-lifecycle-review-profile": _plan_lifecycle_review_profile,
}

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

    cmd = sys.argv[1]
    fn = _COMMANDS.get(cmd)
    if fn is None:
        print(f"Unknown command: {cmd}")
        sys.exit(1)

    try:
        fn()
    except Exception as exc:
        print(f"FAIL [{cmd}]: {type(exc).__name__}: {exc}")
        import traceback

        traceback.print_exc()
        sys.exit(1)