shd.build/tools/game-data-capture/src/game_data_capture/domain/state_machines.py

"""State vocabularies and allowed transitions for recoverable workflows."""

from collections.abc import Mapping
from dataclasses import dataclass
from enum import StrEnum
from types import MappingProxyType
from typing import Final


class CaptureState(StrEnum):
    RECEIVED = "received"
    SPOOLING = "spooling"
    STORED = "stored"
    DEDUPLICATED = "deduplicated"
    QUEUED = "queued"
    FAILED = "failed"
    CANCELLED = "cancelled"


class JobState(StrEnum):
    QUEUED = "queued"
    RUNNING = "running"
    RETRY_WAIT = "retry_wait"
    SUCCEEDED = "succeeded"
    FAILED = "failed"
    CANCELLED = "cancelled"


class ExtractionState(StrEnum):
    PENDING = "pending"
    LOCAL_PROCESSING = "local_processing"
    FALLBACK_PENDING = "fallback_pending"
    FALLBACK_PROCESSING = "fallback_processing"
    CANDIDATE_READY = "candidate_ready"
    ABSTAINED = "abstained"
    FAILED = "failed"
    CANCELLED = "cancelled"


class GroupingState(StrEnum):
    UNGROUPED = "ungrouped"
    PROPOSED = "proposed"
    GROUPED = "grouped"
    SPLIT = "split"
    SUPERSEDED = "superseded"
    REJECTED = "rejected"
    CANCELLED = "cancelled"


class ReviewState(StrEnum):
    PENDING = "pending"
    IN_REVIEW = "in_review"
    CHANGES_REQUESTED = "changes_requested"
    APPROVED = "approved"
    REJECTED = "rejected"
    CANCELLED = "cancelled"


class ExportState(StrEnum):
    NOT_REQUESTED = "not_requested"
    QUEUED = "queued"
    BUILDING = "building"
    READY = "ready"
    FAILED = "failed"
    CANCELLED = "cancelled"
    EXPIRED = "expired"


class ErrorState(StrEnum):
    OPEN = "open"
    RETRYABLE = "retryable"
    RETRYING = "retrying"
    RECOVERED = "recovered"
    TERMINAL = "terminal"
    DISMISSED = "dismissed"


@dataclass(frozen=True, slots=True)
class StateMachine:
    name: str
    transitions: Mapping[str, frozenset[str]]

    def allows(self, current: str, target: str) -> bool:
        return target in self.transitions.get(current, frozenset())


def _machine(name: str, transitions: dict[str, set[str]]) -> StateMachine:
    return StateMachine(
        name=name,
        transitions=MappingProxyType(
            {state: frozenset(targets) for state, targets in transitions.items()}
        ),
    )


STATE_MACHINES: Final[Mapping[str, StateMachine]] = MappingProxyType(
    {
        "capture": _machine(
            "capture",
            {
                "received": {"spooling", "deduplicated", "failed", "cancelled"},
                "spooling": {"stored", "failed", "cancelled"},
                "stored": {"queued", "failed"},
                "deduplicated": set(),
                "queued": set(),
                "failed": {"spooling", "cancelled"},
                "cancelled": set(),
            },
        ),
        "job": _machine(
            "job",
            {
                "queued": {"running", "cancelled"},
                "running": {"succeeded", "retry_wait", "failed", "cancelled"},
                "retry_wait": {"queued", "failed", "cancelled"},
                "succeeded": set(),
                "failed": {"queued"},
                "cancelled": set(),
            },
        ),
        "extraction": _machine(
            "extraction",
            {
                "pending": {"local_processing", "cancelled"},
                "local_processing": {
                    "candidate_ready",
                    "fallback_pending",
                    "abstained",
                    "failed",
                    "cancelled",
                },
                "fallback_pending": {"fallback_processing", "cancelled"},
                "fallback_processing": {
                    "candidate_ready",
                    "abstained",
                    "failed",
                    "cancelled",
                },
                "candidate_ready": set(),
                "abstained": set(),
                "failed": {"pending"},
                "cancelled": set(),
            },
        ),
        "grouping": _machine(
            "grouping",
            {
                "ungrouped": {"proposed", "cancelled"},
                "proposed": {"grouped", "rejected", "cancelled"},
                "grouped": {"split", "superseded"},
                "split": {"superseded"},
                "superseded": set(),
                "rejected": {"ungrouped"},
                "cancelled": set(),
            },
        ),
        "review": _machine(
            "review",
            {
                "pending": {"in_review", "cancelled"},
                "in_review": {
                    "changes_requested",
                    "approved",
                    "rejected",
                    "cancelled",
                },
                "changes_requested": {"in_review", "rejected", "cancelled"},
                "approved": {"in_review"},
                "rejected": {"in_review"},
                "cancelled": set(),
            },
        ),
        "export": _machine(
            "export",
            {
                "not_requested": {"queued"},
                "queued": {"building", "cancelled"},
                "building": {"ready", "failed", "cancelled"},
                "ready": {"expired"},
                "failed": {"queued", "cancelled"},
                "cancelled": set(),
                "expired": {"queued"},
            },
        ),
        "error": _machine(
            "error",
            {
                "open": {"retryable", "terminal", "dismissed"},
                "retryable": {"retrying", "terminal", "dismissed"},
                "retrying": {"recovered", "retryable", "terminal"},
                "recovered": set(),
                "terminal": {"dismissed"},
                "dismissed": set(),
            },
        ),
    }
)


class InvalidStateTransition(ValueError):
    """Raised when a caller attempts a transition absent from the contract."""


def require_transition(machine_name: str, current: str, target: str) -> None:
    machine = STATE_MACHINES.get(machine_name)
    if machine is None:
        raise InvalidStateTransition(f"Unknown state machine: {machine_name}")
    if current not in machine.transitions:
        raise InvalidStateTransition(f"Unknown {machine_name} state: {current}")
    if not machine.allows(current, target):
        raise InvalidStateTransition(f"Invalid {machine_name} transition: {current} -> {target}")