cleveragents-core/robot/e2e/wf04_snapshot_helper.py

"""Snapshot helper for WF04 E2E assertions.

This script is executed by ``wf04_multi_project.robot`` to collect
deterministic parent/subplan metadata from the lifecycle service.
"""

from __future__ import annotations

import json
import sys
import traceback
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[2]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.append(str(SRC))

from cleveragents.application.container import get_container  # noqa: E402


def _iso(value: Any) -> str:
    """Return ISO timestamp string normalised to UTC, or empty string.

    Ensures all timestamps are in a consistent UTC-aware format so that
    lexicographic string comparison in Robot assertions is reliable.
    """
    if value is None:
        return ""
    if not isinstance(value, datetime):
        return ""
    # Normalise naive datetimes by assuming UTC
    if value.tzinfo is None:
        value = value.replace(tzinfo=UTC)
    # Convert to UTC and produce a consistent representation
    return value.astimezone(UTC).isoformat()


def _enum_value(value: Any) -> str:
    """Return enum ``.value`` when available, else ``str(value)``."""
    return value.value if hasattr(value, "value") else str(value)


DecisionTree = dict[str, Any] | list[Any]


def count_decision_nodes(root: DecisionTree | Any, *, max_depth: int = 50) -> int:
    """Recursively count decision nodes in a plan tree structure.

    Walks the tree rooted at *root* and counts every ``dict`` node whose
    ``decision_id`` field is non-empty.  Works for both dict (single root)
    and list (forest) inputs.

    A *max_depth* guard prevents unbounded recursion on malformed trees.
    """
    if max_depth <= 0:
        return 0
    if isinstance(root, list):
        return sum(count_decision_nodes(item, max_depth=max_depth) for item in root)
    if isinstance(root, dict):
        has_decision = 1 if str(root.get("decision_id", "")).strip() else 0
        children = root.get("children") or []
        return has_decision + sum(
            count_decision_nodes(child, max_depth=max_depth - 1) for child in children
        )
    return 0


def _build_snapshot(plan_id: str) -> dict[str, Any]:
    """Build a serialisable snapshot for WF04 verification assertions."""
    lifecycle_service = get_container().plan_lifecycle_service()
    plan = lifecycle_service.get_plan(plan_id)
    if plan is None:
        msg = f"plan '{plan_id}' not found"
        raise ValueError(msg)

    project_scopes: list[dict[str, Any]] = []
    resource_to_project: dict[str, str] = {}

    if plan.multi_project_metadata is not None:
        for scope in plan.multi_project_metadata.project_scopes:
            project_scopes.append(
                {
                    "project_name": scope.project_name,
                    "resource_ids": list(scope.resource_ids),
                }
            )
            for resource_id in scope.resource_ids:
                resource_to_project[resource_id] = scope.project_name

    subplans: list[dict[str, Any]] = []
    for status in plan.subplan_statuses:
        child_plan = lifecycle_service.get_plan(status.subplan_id)
        mapped_projects = sorted(
            {
                resource_to_project[resource_id]
                for resource_id in status.target_resources
                if resource_id in resource_to_project
            }
        )
        unmapped_resources = sorted(
            {
                resource_id
                for resource_id in status.target_resources
                if resource_id not in resource_to_project
            }
        )

        subplans.append(
            {
                "subplan_id": status.subplan_id,
                "status": _enum_value(status.status),
                "started_at": _iso(status.started_at),
                "completed_at": _iso(status.completed_at),
                "target_resources": list(status.target_resources),
                "mapped_projects": mapped_projects,
                "unmapped_resources": unmapped_resources,
                "files_changed": status.files_changed,
                "child_phase": (
                    _enum_value(child_plan.phase) if child_plan is not None else ""
                ),
                "child_state": (
                    _enum_value(child_plan.processing_state)
                    if child_plan is not None
                    else ""
                ),
                "child_updated_at": (
                    _iso(child_plan.timestamps.updated_at)
                    if child_plan is not None
                    else ""
                ),
                "execute_started_at": (
                    _iso(child_plan.timestamps.execute_started_at)
                    if child_plan is not None
                    else ""
                ),
                "execute_completed_at": (
                    _iso(child_plan.timestamps.execute_completed_at)
                    if child_plan is not None
                    else ""
                ),
                "apply_started_at": (
                    _iso(child_plan.timestamps.apply_started_at)
                    if child_plan is not None
                    else ""
                ),
                "applied_at": (
                    _iso(child_plan.timestamps.applied_at)
                    if child_plan is not None
                    else ""
                ),
                "child_validation_summary": (
                    child_plan.validation_summary if child_plan is not None else None
                ),
            }
        )

    return {
        "plan_id": plan.identity.plan_id,
        "subplan_count": len(plan.subplan_statuses),
        "validation_summary": plan.validation_summary,
        "project_scopes": project_scopes,
        "subplans": subplans,
    }


def main() -> int:
    """CLI entrypoint.

    Modes:
        wf04_snapshot_helper.py <plan_id>
            Build and print the full snapshot JSON for a plan.
        wf04_snapshot_helper.py --count-nodes <json_file>
            Read a JSON tree from *json_file* (or ``-`` for stdin) and
            print the decision-node count to stdout.
    """
    if len(sys.argv) == 3 and sys.argv[1] == "--count-nodes":
        return _cli_count_nodes(sys.argv[2])

    if len(sys.argv) != 2:
        print(
            "Usage: wf04_snapshot_helper.py <plan_id>\n"
            "       wf04_snapshot_helper.py --count-nodes <json_file|->\n",
            file=sys.stderr,
        )
        return 2

    plan_id = sys.argv[1]
    try:
        snapshot = _build_snapshot(plan_id)
    except (ValueError, RuntimeError) as exc:
        print(str(exc), file=sys.stderr)
        return 1
    # Intentional catch-all: CLI entrypoint must never propagate
    # unhandled exceptions — report full traceback and exit non-zero.
    except Exception as exc:
        print(f"{exc}\n{traceback.format_exc()}", file=sys.stderr)
        return 1

    print(json.dumps(snapshot, default=str))
    return 0


def _cli_count_nodes(json_source: str) -> int:
    """Read a JSON tree and print the decision-node count."""
    try:
        if json_source == "-":
            tree = json.load(sys.stdin)
        else:
            tree = json.loads(Path(json_source).read_text(encoding="utf-8"))
        print(count_decision_nodes(tree))
    except (ValueError, OSError, json.JSONDecodeError) as exc:
        print(str(exc), file=sys.stderr)
        return 1
    # Intentional catch-all: CLI must not propagate exceptions.
    except Exception as exc:
        print(f"{exc}\n{traceback.format_exc()}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())