pcb-afm-simulation-system/web/backend/app/services/task_contract.py

"""Unified task state vocabulary + adaptive-batch field contract (P3-M4).

Reconciles the two scheduling vocabularies that existed in parallel:

  - TaskManager / Task ORM:
        pending -> dispatched -> running -> completed / failed / cancelled
  - BatchScheduler:
        queued  -> running -> completed / failed / cancelled

This module is the single source of truth for status normalization and the
adaptive-batch fields shared across TaskManager, BatchScheduler and the
adaptive orchestrator. It must stay dependency-free so any of those services
can import it without cycles.

All source is ASCII only.
"""
from typing import Any, Dict, Optional

# Canonical statuses (TaskManager vocabulary).
TASK_STATUS_PENDING = "pending"
TASK_STATUS_DISPATCHED = "dispatched"
TASK_STATUS_RUNNING = "running"
TASK_STATUS_COMPLETED = "completed"
TASK_STATUS_COMPLETED_WITH_ERRORS = "completed_with_errors"
TASK_STATUS_FAILED = "failed"
TASK_STATUS_CANCELLED = "cancelled"

# BatchScheduler-only vocabulary.
TASK_STATUS_QUEUED = "queued"

TERMINAL_STATUSES = {
    TASK_STATUS_COMPLETED,
    TASK_STATUS_COMPLETED_WITH_ERRORS,
    TASK_STATUS_FAILED,
    TASK_STATUS_CANCELLED,
}

# scheduler / legacy word -> canonical word
STATUS_ALIASES = {
    TASK_STATUS_QUEUED: TASK_STATUS_PENDING,
    TASK_STATUS_PENDING: TASK_STATUS_PENDING,
    TASK_STATUS_DISPATCHED: TASK_STATUS_DISPATCHED,
    TASK_STATUS_RUNNING: TASK_STATUS_RUNNING,
    TASK_STATUS_COMPLETED: TASK_STATUS_COMPLETED,
    TASK_STATUS_COMPLETED_WITH_ERRORS: TASK_STATUS_COMPLETED,
    TASK_STATUS_FAILED: TASK_STATUS_FAILED,
    TASK_STATUS_CANCELLED: TASK_STATUS_CANCELLED,
    "canceled": TASK_STATUS_CANCELLED,
}

# Adaptive-batch fields introduced in P3-M2, shared by Task ORM,
# BatchScheduler and AdaptiveOrchestrator.
ADAPTIVE_BATCH_FIELDS = ("task_type", "loop_id", "batch_id", "point_ids", "dynamic")


def normalize_status(status: Optional[str]) -> Optional[str]:
    """Map any scheduler/legacy status word to the canonical status.

    Unknown words pass through unchanged so validation can report them.
    """
    if status is None:
        return None
    return STATUS_ALIASES.get(str(status), str(status))


def is_terminal(status: Optional[str]) -> bool:
    return normalize_status(status) in TERMINAL_STATUSES


def merge_adaptive_fields(task_info: Dict[str, Any], **kwargs) -> Dict[str, Any]:
    """Merge adaptive-batch fields into a task dict (no-ops stay default).

    Accepts kwargs that may include task_type/loop_id/batch_id/point_ids/
    dynamic and sets the canonical defaults when absent.
    """
    out = dict(task_info)
    out.setdefault("task_type", kwargs.get("task_type", "scan"))
    out.setdefault("loop_id", kwargs.get("loop_id"))
    out.setdefault("batch_id", kwargs.get("batch_id"))
    out.setdefault("point_ids", kwargs.get("point_ids") or [])
    out.setdefault("dynamic", bool(kwargs.get("dynamic", False)))
    return out