stonks-oracle/services/intelligence_pipeline_v3/orchestrator/state.py

"""Pipeline and stage state machines with explicit transitions and idempotency."""

from __future__ import annotations

import enum
import hashlib
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from uuid import UUID, uuid4


class PipelineState(str, enum.Enum):
    """Top-level pipeline run states."""

    PENDING = "pending"
    SEGMENTING = "segmenting"
    EXTRACTING = "extracting"
    RESOLVING = "resolving"
    VERIFYING = "verifying"
    ROUTING = "routing"
    ADJUDICATING = "adjudicating"
    IMPACT = "impact"
    PERSISTING = "persisting"
    COMPLETED = "completed"
    FAILED = "failed"
    DEAD_LETTER = "dead_letter"


class StageState(str, enum.Enum):
    """Per-stage execution states."""

    QUEUED = "queued"
    LEASED = "leased"
    RUNNING = "running"
    SUCCEEDED = "succeeded"
    RETRYING = "retrying"
    FAILED = "failed"
    SKIPPED = "skipped"


# Valid transitions for the pipeline state machine
_PIPELINE_TRANSITIONS: dict[PipelineState, set[PipelineState]] = {
    PipelineState.PENDING: {PipelineState.SEGMENTING, PipelineState.FAILED},
    PipelineState.SEGMENTING: {PipelineState.EXTRACTING, PipelineState.FAILED},
    PipelineState.EXTRACTING: {PipelineState.RESOLVING, PipelineState.FAILED},
    PipelineState.RESOLVING: {PipelineState.VERIFYING, PipelineState.FAILED},
    PipelineState.VERIFYING: {PipelineState.ROUTING, PipelineState.FAILED},
    PipelineState.ROUTING: {
        PipelineState.ADJUDICATING,
        PipelineState.IMPACT,
        PipelineState.FAILED,
    },
    PipelineState.ADJUDICATING: {PipelineState.IMPACT, PipelineState.FAILED},
    PipelineState.IMPACT: {PipelineState.PERSISTING, PipelineState.FAILED},
    PipelineState.PERSISTING: {PipelineState.COMPLETED, PipelineState.FAILED},
    PipelineState.COMPLETED: set(),
    PipelineState.FAILED: {PipelineState.DEAD_LETTER, PipelineState.PENDING},
    PipelineState.DEAD_LETTER: set(),
}

# Valid transitions for stage states
_STAGE_TRANSITIONS: dict[StageState, set[StageState]] = {
    StageState.QUEUED: {StageState.LEASED, StageState.SKIPPED},
    StageState.LEASED: {StageState.RUNNING, StageState.QUEUED},
    StageState.RUNNING: {StageState.SUCCEEDED, StageState.RETRYING, StageState.FAILED},
    StageState.SUCCEEDED: set(),
    StageState.RETRYING: {StageState.QUEUED, StageState.FAILED},
    StageState.FAILED: set(),
    StageState.SKIPPED: set(),
}


@dataclass(frozen=True)
class StateTransition:
    """Immutable record of a state transition."""

    transition_id: UUID
    run_id: UUID
    from_state: PipelineState | StageState
    to_state: PipelineState | StageState
    timestamp: datetime
    reason: str
    idempotency_key: str


def _compute_idempotency_key(
    document_id: str, stage: str, attempt: int
) -> str:
    """Deterministic idempotency key from document, stage, and attempt."""
    raw = f"{document_id}:{stage}:{attempt}"
    return hashlib.sha256(raw.encode()).hexdigest()[:32]


@dataclass
class PipelineStateMachine:
    """Manages state transitions for a single pipeline run.

    Enforces valid transitions, records history, and generates
    idempotency keys for each stage attempt.
    """

    run_id: UUID = field(default_factory=uuid4)
    document_id: str = ""
    state: PipelineState = PipelineState.PENDING
    stage_states: dict[str, StageState] = field(default_factory=dict)
    stage_attempts: dict[str, int] = field(default_factory=dict)
    history: list[StateTransition] = field(default_factory=list)
    max_retries: int = 3
    metadata: dict[str, Any] = field(default_factory=dict)

    def transition_pipeline(
        self, to_state: PipelineState, reason: str = ""
    ) -> StateTransition:
        """Advance the pipeline to a new state.

        Raises ValueError if the transition is invalid.
        """
        allowed = _PIPELINE_TRANSITIONS.get(self.state, set())
        if to_state not in allowed:
            raise ValueError(
                f"Invalid pipeline transition: {self.state.value} -> {to_state.value}"
            )

        transition = StateTransition(
            transition_id=uuid4(),
            run_id=self.run_id,
            from_state=self.state,
            to_state=to_state,
            timestamp=datetime.now(timezone.utc),
            reason=reason,
            idempotency_key=_compute_idempotency_key(
                self.document_id, to_state.value, 0
            ),
        )
        self.state = to_state
        self.history.append(transition)
        return transition

    def transition_stage(
        self, stage: str, to_state: StageState, reason: str = ""
    ) -> StateTransition:
        """Advance a stage to a new state.

        Raises ValueError if the transition is invalid.
        """
        current = self.stage_states.get(stage, StageState.QUEUED)
        allowed = _STAGE_TRANSITIONS.get(current, set())
        if to_state not in allowed:
            raise ValueError(
                f"Invalid stage transition for '{stage}': "
                f"{current.value} -> {to_state.value}"
            )

        attempt = self.stage_attempts.get(stage, 0)
        if to_state == StageState.RETRYING:
            attempt += 1
            self.stage_attempts[stage] = attempt

        transition = StateTransition(
            transition_id=uuid4(),
            run_id=self.run_id,
            from_state=current,
            to_state=to_state,
            timestamp=datetime.now(timezone.utc),
            reason=reason,
            idempotency_key=_compute_idempotency_key(
                self.document_id, stage, attempt
            ),
        )
        self.stage_states[stage] = to_state
        self.history.append(transition)
        return transition

    def can_retry(self, stage: str) -> bool:
        """Check whether the stage has retries remaining."""
        return self.stage_attempts.get(stage, 0) < self.max_retries

    def should_dead_letter(self) -> bool:
        """Check if the pipeline run should move to dead letter."""
        if self.state != PipelineState.FAILED:
            return False
        # Dead-letter if any stage exceeded max retries
        for stage, attempts in self.stage_attempts.items():
            if attempts >= self.max_retries:
                return True
        return False