Source code for rath.runtime.models

"""Durable Run, Checkpoint, Interrupt, and event state models."""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from enum import Enum
from uuid import UUID, uuid4

from rath._json import JSONValue, freeze_mapping
from rath.errors import ErrorCode, RathError

__all__ = [
    "ApprovalDecision",
    "ApprovalDecisionKind",
    "Checkpoint",
    "ConflictError",
    "Interrupt",
    "InterruptKind",
    "InvalidRunTransition",
    "Run",
    "RunEvent",
    "RunStatus",
    "ClaimedRun",
    "ResourceLease",
    "assert_transition",
]


[docs] class RunStatus(str, Enum): QUEUED = "queued" RUNNING = "running" WAITING = "waiting" SUCCEEDED = "succeeded" FAILED = "failed" CANCELLED = "cancelled" TIMED_OUT = "timed_out" NEEDS_REVIEW = "needs_review"
TERMINAL_RUN_STATUSES = frozenset( { RunStatus.SUCCEEDED, RunStatus.FAILED, RunStatus.CANCELLED, RunStatus.TIMED_OUT, } ) _TRANSITIONS: Mapping[RunStatus, frozenset[RunStatus]] = { RunStatus.QUEUED: frozenset( { RunStatus.RUNNING, RunStatus.CANCELLED, RunStatus.TIMED_OUT, } ), RunStatus.RUNNING: frozenset( { RunStatus.QUEUED, RunStatus.WAITING, RunStatus.SUCCEEDED, RunStatus.FAILED, RunStatus.CANCELLED, RunStatus.TIMED_OUT, RunStatus.NEEDS_REVIEW, } ), RunStatus.WAITING: frozenset( { RunStatus.QUEUED, RunStatus.CANCELLED, RunStatus.TIMED_OUT, RunStatus.FAILED, } ), RunStatus.NEEDS_REVIEW: frozenset( { RunStatus.QUEUED, RunStatus.FAILED, RunStatus.CANCELLED, } ), RunStatus.SUCCEEDED: frozenset(), RunStatus.FAILED: frozenset(), RunStatus.CANCELLED: frozenset(), RunStatus.TIMED_OUT: frozenset(), }
[docs] class ConflictError(RathError): def __init__(self, message: str, *, details: Mapping[str, object] | None = None): super().__init__( ErrorCode.CONFLICT, message, retryable=False, details=details, )
[docs] class InvalidRunTransition(ConflictError): def __init__(self, source: RunStatus, target: RunStatus) -> None: super().__init__( f"invalid run transition from {source.value!r} to {target.value!r}", details={"source": source.value, "target": target.value}, ) self.source = source self.target = target
[docs] def assert_transition(source: RunStatus, target: RunStatus) -> None: if target not in _TRANSITIONS[source]: raise InvalidRunTransition(source, target)
def _aware(value: datetime, *, field_name: str) -> None: if value.tzinfo is None: raise ValueError(f"{field_name} must be timezone-aware")
[docs] @dataclass(frozen=True, slots=True) class Run: id: UUID plan_id: UUID revision_id: UUID session_id: UUID tenant_id: str status: RunStatus state: Mapping[str, JSONValue] next_nodes: tuple[str, ...] created_at: datetime updated_at: datetime version: int = 0 idempotency_key: str | None = None context: Mapping[str, JSONValue] = field(default_factory=dict) priority: int = 0 def __post_init__(self) -> None: if not self.tenant_id.strip(): raise ValueError("run tenant_id must not be empty") if self.version < 0: raise ValueError("run version must not be negative") if not -100 <= self.priority <= 100: raise ValueError("run priority must be between -100 and 100") _aware(self.created_at, field_name="run.created_at") _aware(self.updated_at, field_name="run.updated_at") object.__setattr__(self, "state", freeze_mapping(self.state, field="run.state")) object.__setattr__( self, "context", freeze_mapping(self.context, field="run.context") ) object.__setattr__(self, "next_nodes", tuple(self.next_nodes))
[docs] @classmethod def create( cls, *, plan_id: UUID, revision_id: UUID, session_id: UUID, tenant_id: str, status: RunStatus = RunStatus.QUEUED, state: Mapping[str, object] | None = None, next_nodes: tuple[str, ...] = (), idempotency_key: str | None = None, context: Mapping[str, object] | None = None, priority: int = 0, id: UUID | None = None, ) -> "Run": now = datetime.now(timezone.utc) return cls( id=id or uuid4(), plan_id=plan_id, revision_id=revision_id, session_id=session_id, tenant_id=tenant_id, status=status, state=freeze_mapping(state, field="run.state"), next_nodes=next_nodes, idempotency_key=idempotency_key, context=freeze_mapping(context, field="run.context"), priority=priority, created_at=now, updated_at=now, )
[docs] @dataclass(frozen=True, slots=True) class RunEvent: run_id: UUID sequence: int type: str data: Mapping[str, JSONValue] created_at: datetime def __post_init__(self) -> None: if self.sequence < 1: raise ValueError("run event sequence must be greater than zero") _aware(self.created_at, field_name="run_event.created_at") object.__setattr__( self, "data", freeze_mapping(self.data, field="run_event.data"), )
[docs] @dataclass(frozen=True, slots=True) class Checkpoint: id: UUID run_id: UUID sequence: int plan_hash: str state: Mapping[str, JSONValue] next_nodes: tuple[str, ...] pending_interrupts: tuple[UUID, ...] effect_watermark: int created_at: datetime def __post_init__(self) -> None: if self.sequence < 1: raise ValueError("checkpoint sequence must be greater than zero") if self.effect_watermark < 0: raise ValueError("effect_watermark must not be negative") if len(self.plan_hash) != 64: raise ValueError("plan_hash must be a SHA-256 hexadecimal digest") try: int(self.plan_hash, 16) except ValueError as exc: raise ValueError("plan_hash must be hexadecimal") from exc _aware(self.created_at, field_name="checkpoint.created_at") object.__setattr__( self, "state", freeze_mapping(self.state, field="checkpoint.state"), ) object.__setattr__(self, "next_nodes", tuple(self.next_nodes)) object.__setattr__( self, "pending_interrupts", tuple(self.pending_interrupts), )
[docs] @classmethod def create( cls, *, run_id: UUID, sequence: int, plan_hash: str, state: Mapping[str, object], next_nodes: tuple[str, ...], effect_watermark: int, pending_interrupts: tuple[UUID, ...] = (), ) -> "Checkpoint": return cls( id=uuid4(), run_id=run_id, sequence=sequence, plan_hash=plan_hash, state=freeze_mapping(state, field="checkpoint.state"), next_nodes=next_nodes, pending_interrupts=pending_interrupts, effect_watermark=effect_watermark, created_at=datetime.now(timezone.utc), )
[docs] class InterruptKind(str, Enum): APPROVAL = "approval" INPUT = "input" REVIEW = "review"
[docs] class ApprovalDecisionKind(str, Enum): APPROVE = "approve" EDIT = "edit" REJECT = "reject" RESPOND = "respond"
[docs] @dataclass(frozen=True, slots=True) class ApprovalDecision: kind: ApprovalDecisionKind actor_id: str reason: str payload: Mapping[str, JSONValue] = field(default_factory=dict) def __post_init__(self) -> None: if not self.actor_id.strip(): raise ValueError("decision actor_id must not be empty") if not self.reason.strip(): raise ValueError("decision reason must not be empty") object.__setattr__( self, "payload", freeze_mapping(self.payload, field="decision.payload"), )
[docs] @dataclass(frozen=True, slots=True) class Interrupt: id: UUID run_id: UUID kind: InterruptKind request: Mapping[str, JSONValue] created_at: datetime expires_at: datetime | None = None decision: ApprovalDecision | None = None decided_at: datetime | None = None def __post_init__(self) -> None: _aware(self.created_at, field_name="interrupt.created_at") if self.expires_at is not None: _aware(self.expires_at, field_name="interrupt.expires_at") if self.expires_at <= self.created_at: raise ValueError("interrupt expires_at must be after created_at") if self.decided_at is not None: _aware(self.decided_at, field_name="interrupt.decided_at") if (self.decision is None) != (self.decided_at is None): raise ValueError("decision and decided_at must be set together") object.__setattr__( self, "request", freeze_mapping(self.request, field="interrupt.request"), )
[docs] @classmethod def create( cls, *, run_id: UUID, kind: InterruptKind, request: Mapping[str, object], timeout_seconds: float | None = None, ) -> "Interrupt": if timeout_seconds is not None and timeout_seconds <= 0: raise ValueError("interrupt timeout_seconds must be positive") created_at = datetime.now(timezone.utc) return cls( id=uuid4(), run_id=run_id, kind=kind, request=freeze_mapping(request, field="interrupt.request"), created_at=created_at, expires_at=( created_at + timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None ), )
[docs] @dataclass(frozen=True, slots=True) class ResourceLease: id: UUID resource_type: str resource_id: str owner_run_id: UUID holder_worker_id: str expires_at: datetime fencing_token: int created_at: datetime updated_at: datetime def __post_init__(self) -> None: if not self.resource_type.strip(): raise ValueError("lease resource_type must not be empty") if not self.resource_id.strip(): raise ValueError("lease resource_id must not be empty") if not self.holder_worker_id.strip(): raise ValueError("lease holder_worker_id must not be empty") if self.fencing_token < 1: raise ValueError("lease fencing_token must be positive") _aware(self.expires_at, field_name="lease.expires_at") _aware(self.created_at, field_name="lease.created_at") _aware(self.updated_at, field_name="lease.updated_at")
[docs] @dataclass(frozen=True, slots=True) class ClaimedRun: run: Run lease: ResourceLease