src/progress.py

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

# Exact worker-facing steps emitted by src/worker_payload.py.
# Keep this list close to the Job logs/events rather than a marketing pipeline.
STEP_ORDER = [
    "bootstrap",
    "dependencies",
    "auth",
    "model_analysis",
    "workspace",
    "node",
    "pi_install",
    "pi_config",
    "pi_run",
    "pi_model_resolution",
    "pi_verification",
    "metadata_sanitize",
    "requirements_sanitize",
    "hardware_strategy",
    "create_space_hardware",
    "create_space",
    "failure_detected",
    "failure_diagnosis",
    "pi_diagnosis",
    "repair_decision",
    "wait_for_logs",
    "factory_rebuild",
    "repair",
    "repair_diagnosis",
    "repair_brief",
    "repair_plan",
    "repair_patch",
    "repair_upload",
    "repair_validation",
    "technical_blocker",
    "manual_hardware_required",
    "upload_files",
    "space_runtime",
    "space_logs",
    "api_validation",
    "live_wait",
    "generation_smoke",
    "inference_gate",
    "report_write",
    "done",
    "failure",
]

STEP_LABELS = {
    "bootstrap": "Bootstrap",
    "dependencies": "Dependencies",
    "auth": "Auth",
    "model_analysis": "Model analysis",
    "workspace": "Workspace",
    "node": "Node/npm",
    "pi_install": "Pi install",
    "pi_config": "Pi config",
    "pi_run": "Pi run",
    "pi_model_resolution": "Pi model",
    "pi_verification": "Pi verification",
    "metadata_sanitize": "Metadata sanitize",
    "requirements_sanitize": "Requirements sanitize",
    "hardware_strategy": "Hardware strategy",
    "create_space_hardware": "Create with hardware",
    "create_space": "Create Space",
    "failure_detected": "Failure detected",
    "failure_diagnosis": "Failure diagnosis",
    "pi_diagnosis": "Pi diagnosis",
    "repair_decision": "Repair decision",
    "wait_for_logs": "Wait for logs",
    "factory_rebuild": "Factory rebuild",
    "repair": "Repair start",
    "repair_diagnosis": "Patch allowed",
    "repair_brief": "Repair brief",
    "repair_plan": "Repair plan",
    "repair_patch": "Repair patch",
    "repair_upload": "Repair upload",
    "repair_validation": "Repair validation",
    "technical_blocker": "Technical blocker",
    "manual_hardware_required": "Manual hardware",
    "upload_files": "Upload files",
    "space_runtime": "Space runtime",
    "space_logs": "Space logs",
    "api_validation": "API validation",
    "live_wait": "Live wait",
    "generation_smoke": "Generation smoke",
    "inference_gate": "Inference gate",
    "report_write": "Report",
    "done": "Done",
    "failure": "Failure",
}

STEP_ALIASES = {
    "bucket_ready": "bootstrap",
    "job_launched": "bootstrap",
    "hardware_preferred": "hardware_strategy",
    "hardware_fallback": "hardware_strategy",
    "hardware": "hardware_strategy",
    "build": "space_runtime",
    "diagnose": "failure_diagnosis",
    "diagnosis": "failure_diagnosis",
    "pi_decision": "repair_decision",
    "decision": "repair_decision",
    "patch": "repair_patch",
    "repair_failed": "repair_validation",
}

DONE_STATUSES = {
    "success",
    "done",
    "completed",
    "passed",
    "full_inference_success",
    "manual_hardware_required",
    "technical_blocker",
}
PARTIAL_STATUSES = {
    "full_inference_candidate_health_passed",
    "health_only",
    "partial",
    "partial_validation",
    "completed_with_warnings",
}
RUNNING_STATUSES = {"started", "running", "waiting", "pending", "scheduled"}
FAILED_STATUSES = {"failed", "error", "failure"}
CANCELLED_STATUSES = {"cancelled", "canceled"}


def _parse_ts(ts: str | None) -> datetime | None:
    if not ts:
        return None
    try:
        return datetime.fromisoformat(ts.replace("Z", "+00:00"))
    except Exception:
        return None


def _canonical_step(step: str | None) -> str | None:
    if not step:
        return None
    value = str(step)
    return STEP_ALIASES.get(value, value if value in STEP_ORDER else None)


def _is_terminal_status(status: str | None) -> bool:
    s = str(status or "").lower()
    return s in DONE_STATUSES or s in PARTIAL_STATUSES or s in FAILED_STATUSES or s in CANCELLED_STATUSES


def _progress_for_index(index: int) -> int:
    if index <= 0:
        return 3
    if len(STEP_ORDER) <= 1:
        return 0
    return int(round((index / (len(STEP_ORDER) - 1)) * 100))


def progress_from_events(events: list[dict[str, Any]] | None, *, state: dict[str, Any] | None = None) -> dict[str, Any]:
    """Build a monotonic live progress model from worker events + Job state.

    The UI was confusing when late/partial bucket events made steps appear to
    move backwards. This model treats the furthest observed real worker step as
    the source of truth, marks all earlier steps done, keeps only that step
    running while the Job is active, and freezes it as stopped on terminal
    non-success statuses.
    """
    events = events or []
    state = state or {}
    step_status = {step: "pending" for step in STEP_ORDER}
    last_event: dict[str, Any] | None = None
    terminal_status = None
    first_ts = None
    last_ts = None
    furthest_index = -1
    current_step = "bootstrap"
    current_status = ""
    current_event: dict[str, Any] | None = None
    latest_signal_step = ""
    latest_signal_status = ""
    latest_signal_event: dict[str, Any] | None = None
    latest_signal_ts: datetime | None = None

    for event in events:
        if not isinstance(event, dict):
            continue
        step = _canonical_step(event.get("step"))
        status = str(event.get("status") or "").lower()
        ts = _parse_ts(event.get("ts"))
        if ts and not first_ts:
            first_ts = ts
        if ts:
            last_ts = ts
        if not step:
            continue

        last_event = event
        idx = STEP_ORDER.index(step)
        # Never let older / lower-index failed events poison the whole run once
        # later validation or repair steps have appeared.
        if idx >= furthest_index:
            furthest_index = idx
            current_step = step
            current_status = status
            current_event = event
        if ts is None or latest_signal_ts is None or ts >= latest_signal_ts:
            latest_signal_ts = ts
            latest_signal_step = step
            latest_signal_status = status
            latest_signal_event = event

    if latest_signal_status in RUNNING_STATUSES:
        current_step = latest_signal_step or current_step
        current_status = latest_signal_status
        current_event = latest_signal_event or current_event
        if current_step in STEP_ORDER:
            furthest_index = STEP_ORDER.index(current_step)
        terminal_status = None

    if current_step == "done" and current_status:
        terminal_status = current_status
    elif current_step == "failure" and current_status:
        terminal_status = "failed" if current_status in FAILED_STATUSES or current_status in DONE_STATUSES else current_status
    elif current_status in CANCELLED_STATUSES:
        terminal_status = "cancelled"
    elif current_status in FAILED_STATUSES:
        terminal_status = "failed"

    status_from_state = str(state.get("status") or state.get("gate_status") or "").lower()
    job_stage = str(state.get("job_stage") or "").lower()
    if status_from_state in CANCELLED_STATUSES:
        terminal_status = "cancelled"
    elif status_from_state in PARTIAL_STATUSES:
        terminal_status = status_from_state
    elif status_from_state in DONE_STATUSES:
        terminal_status = status_from_state
    elif status_from_state in FAILED_STATUSES:
        terminal_status = "failed"
        if current_status not in FAILED_STATUSES and current_step != "failure":
            current_step = "failure"
            current_status = "failed"
            furthest_index = STEP_ORDER.index("failure")
    elif job_stage in CANCELLED_STATUSES:
        terminal_status = "cancelled"
    elif job_stage in FAILED_STATUSES:
        # A terminal HF Job failure must produce an explicit red timeline point.
        terminal_status = "failed"
        if current_status not in FAILED_STATUSES and current_step != "failure":
            current_step = "failure"
            current_status = "failed"
            furthest_index = STEP_ORDER.index("failure")

    if terminal_status:
        overall_status = terminal_status
    elif current_status in RUNNING_STATUSES:
        overall_status = "running"
    elif status_from_state in RUNNING_STATUSES:
        overall_status = "running"
    elif events:
        overall_status = "running"
    else:
        overall_status = "not_started"
    terminal = _is_terminal_status(overall_status)

    if furthest_index < 0:
        furthest_index = 0

    if overall_status in DONE_STATUSES or overall_status in PARTIAL_STATUSES or current_step == "done":
        current_step = "done"
        furthest_index = STEP_ORDER.index("done")
        for step in STEP_ORDER:
            if step != "failure":
                step_status[step] = "done"
    elif overall_status in FAILED_STATUSES:
        failure_idx = STEP_ORDER.index("failure")
        if current_step in STEP_ORDER and current_step != "failure":
            # Preserve the failed worker step when it is known; otherwise expose
            # the explicit Failure stage so the UI always has a red dot.
            failed_idx = STEP_ORDER.index(current_step)
        else:
            failed_idx = failure_idx
            current_step = "failure"
        furthest_index = max(furthest_index, failed_idx)
        for idx, step in enumerate(STEP_ORDER):
            if idx < failed_idx:
                step_status[step] = "done"
            elif idx == failed_idx:
                step_status[step] = "failed"
            else:
                step_status[step] = "pending"
        step_status["failure"] = "failed"
    else:
        for idx, step in enumerate(STEP_ORDER):
            if idx < furthest_index:
                step_status[step] = "done"
            elif idx == furthest_index:
                if overall_status in FAILED_STATUSES:
                    step_status[step] = "failed"
                elif overall_status in CANCELLED_STATUSES or (terminal and overall_status not in DONE_STATUSES):
                    step_status[step] = "stopped"
                elif overall_status in {"not_started", "idle"}:
                    step_status[step] = "pending"
                else:
                    step_status[step] = "running"
            else:
                step_status[step] = "pending"

    if overall_status in FAILED_STATUSES:
        progress = max(_progress_for_index(furthest_index), 3)
    elif overall_status in CANCELLED_STATUSES:
        progress = max(_progress_for_index(furthest_index), 3)
    elif overall_status in DONE_STATUSES or overall_status in PARTIAL_STATUSES or current_step == "done":
        progress = 100
    else:
        progress = max(3, min(99, _progress_for_index(furthest_index)))

    now = datetime.now(timezone.utc)
    started = first_ts or _parse_ts(str(state.get("started_at") or state.get("created_at") or ""))
    if started:
        end = (last_ts if terminal and last_ts else now)
        elapsed = int((end - started).total_seconds())
    else:
        elapsed = 0

    timeline = [
        {"step": step, "label": STEP_LABELS[step], "status": step_status[step]}
        for step in STEP_ORDER
    ]

    last_message = None
    if current_event:
        last_message = current_event.get("message") or current_event.get("step")
    elif last_event:
        last_message = last_event.get("message") or last_event.get("step")

    return {
        "status": overall_status,
        "progress": int(max(0, min(100, progress))),
        "current_step": current_step,
        "current_step_label": STEP_LABELS.get(current_step, current_step),
        "last_event": last_message or "No events yet",
        "visual_status": "error" if overall_status in FAILED_STATUSES else ("warn" if overall_status in PARTIAL_STATUSES else ("stopped" if overall_status in CANCELLED_STATUSES else ("success" if overall_status in DONE_STATUSES and overall_status not in {"manual_hardware_required", "technical_blocker"} else ("stopped" if overall_status in {"manual_hardware_required", "technical_blocker"} else ("running" if overall_status == "running" else "neutral"))))),
        "elapsed_seconds": max(0, elapsed),
        "eta_seconds": None,
        "timeline": timeline,
        "last_event_raw": last_event or {},
    }