app/generation/domain/state_machine.py

from __future__ import annotations

from app.generation.domain.enums import GenerationJobStatus
from app.generation.domain.errors import GenerationTransitionError


ALLOWED_TRANSITIONS: dict[GenerationJobStatus, frozenset[GenerationJobStatus]] = {
    GenerationJobStatus.QUEUED: frozenset(
        {
            GenerationJobStatus.SUBMITTING,
            # A trusted remote worker can acknowledge a queued/running or
            # already-completed job before the next reconciliation pass.  The
            # repository only permits these transitions after an opaque worker
            # job ID has been bound; public callers cannot make them.
            GenerationJobStatus.RUNNING,
            GenerationJobStatus.SUCCEEDED,
            GenerationJobStatus.CANCEL_REQUESTED,
            GenerationJobStatus.CANCELLED,
            GenerationJobStatus.FAILED,
        }
    ),
    GenerationJobStatus.SUBMITTING: frozenset(
        {
            GenerationJobStatus.RUNNING,
            GenerationJobStatus.QUEUED,
            GenerationJobStatus.SUCCEEDED,
            GenerationJobStatus.RETRYING,
            GenerationJobStatus.FAILED,
            GenerationJobStatus.CANCEL_REQUESTED,
            GenerationJobStatus.CANCELLED,
        }
    ),
    GenerationJobStatus.RUNNING: frozenset(
        {
            GenerationJobStatus.SUCCEEDED,
            GenerationJobStatus.RETRYING,
            GenerationJobStatus.FAILED,
            GenerationJobStatus.CANCEL_REQUESTED,
            GenerationJobStatus.CANCELLED,
        }
    ),
    GenerationJobStatus.RETRYING: frozenset(
        {
            GenerationJobStatus.SUBMITTING,
            GenerationJobStatus.FAILED,
            GenerationJobStatus.CANCELLED,
        }
    ),
    GenerationJobStatus.CANCEL_REQUESTED: frozenset(
        {
            GenerationJobStatus.CANCELLED,
            GenerationJobStatus.SUCCEEDED,
            GenerationJobStatus.FAILED,
        }
    ),
    GenerationJobStatus.SUCCEEDED: frozenset(),
    GenerationJobStatus.FAILED: frozenset(),
    GenerationJobStatus.CANCELLED: frozenset(),
}


def validate_transition(
    current: str | GenerationJobStatus, target: str | GenerationJobStatus
) -> GenerationJobStatus:
    source = GenerationJobStatus(current)
    destination = GenerationJobStatus(target)
    if destination not in ALLOWED_TRANSITIONS[source]:
        raise GenerationTransitionError(
            f"Cannot transition generation job from {source.value} to {destination.value}."
        )
    return destination