aios/services/aios-autonomous-release-candidate/aios_autonomous_release_candidate/workflow.py

"""Durable fail-closed workflow state machine."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from .errors import StateConflict
from .models import StepState, TaskPlan, TaskState, is_sha256
from .store import StateStore


@dataclass(frozen=True, slots=True)
class WorkflowDecision:
    task_id: str
    ready_steps: tuple[str, ...]
    blocked: bool
    reason: str | None = None


class WorkflowEngine:
    def __init__(self, store: StateStore) -> None:
        self.store = store

    def _task_state(self, task_id: str, state: TaskState) -> None:
        current, revision = self.store.get_task_state(task_id)
        if current != state:
            self.store.compare_and_set_task_state(task_id, revision, state)

    def plan(self, task: TaskPlan) -> None:
        self.store.create_task(
            task.task_id,
            task.digest(),
            {
                "task_id": task.task_id,
                "intent": task.intent,
                "branch": task.binding.branch,
                "base_oid": task.binding.base_oid,
            },
            tuple(step.step_id for step in task.steps),
            {step.step_id: step.max_attempts for step in task.steps},
        )

    def next_steps(self, task: TaskPlan) -> WorkflowDecision:
        states = self.store.list_steps(task.task_id)
        by_id = {step.step_id: step for step in task.steps}
        if any(item.state == StepState.FAILED for item in states.values()):
            self._task_state(task.task_id, TaskState.FAILED)
            return WorkflowDecision(task.task_id, (), True, "terminal_step_failure")
        ready: list[str] = []
        for step_id, stored in states.items():
            if stored.state not in (StepState.PENDING, StepState.READY):
                continue
            dependencies = by_id[step_id].depends_on
            dependency_states = tuple(states[item].state for item in dependencies)
            if any(item in (StepState.FAILED, StepState.BLOCKED) for item in dependency_states):
                if stored.state != StepState.BLOCKED:
                    self.store.transition_step(
                        task.task_id,
                        step_id,
                        (stored.state,),
                        StepState.BLOCKED,
                        last_error="dependency_failed",
                    )
                continue
            if all(item == StepState.COMPLETED for item in dependency_states):
                if stored.state == StepState.PENDING:
                    self.store.transition_step(
                        task.task_id, step_id, (StepState.PENDING,), StepState.READY
                    )
                ready.append(step_id)
        if states and all(item.state == StepState.COMPLETED for item in states.values()):
            self._task_state(task.task_id, TaskState.COMPLETED)
        elif ready:
            self._task_state(task.task_id, TaskState.READY)
        return WorkflowDecision(task.task_id, tuple(sorted(ready)), False)

    def claim(self, task_id: str, step_id: str, worker_id: str, session_id: str) -> None:
        if not worker_id or not session_id:
            raise StateConflict("STEP_CLAIM_BINDING_REQUIRED", "worker and session are required")
        self.store.transition_step(
            task_id,
            step_id,
            (StepState.READY,),
            StepState.CLAIMED,
            worker_id=worker_id,
            session_id=session_id,
            increment_attempt=True,
        )

    def start(
        self,
        task_id: str,
        step_id: str,
        worker_id: str,
        session_id: str,
        *,
        lease_seconds: int = 60,
        now: datetime | None = None,
    ) -> None:
        current = self.store.list_steps(task_id)[step_id]
        if current.worker_id != worker_id or current.session_id != session_id:
            raise StateConflict("STEP_CLAIM_OWNER_MISMATCH", "claim owner does not match")
        now = now or datetime.now(timezone.utc)
        self.store.transition_step(
            task_id,
            step_id,
            (StepState.CLAIMED,),
            StepState.RUNNING,
            worker_id=worker_id,
            session_id=session_id,
            lease_expires_at=now + timedelta(seconds=lease_seconds),
        )
        self._task_state(task_id, TaskState.RUNNING)

    def heartbeat(
        self,
        task_id: str,
        step_id: str,
        worker_id: str,
        session_id: str,
        *,
        lease_seconds: int = 60,
        now: datetime | None = None,
    ) -> None:
        current = self.store.list_steps(task_id)[step_id]
        if current.worker_id != worker_id or current.session_id != session_id:
            raise StateConflict("STEP_CLAIM_OWNER_MISMATCH", "heartbeat owner does not match")
        now = now or datetime.now(timezone.utc)
        self.store.transition_step(
            task_id,
            step_id,
            (StepState.RUNNING,),
            StepState.RUNNING,
            lease_expires_at=now + timedelta(seconds=lease_seconds),
        )

    def complete(self, task_id: str, step_id: str, *, result_digest: str) -> None:
        if not is_sha256(result_digest):
            raise StateConflict("RESULT_DIGEST_INVALID", "result digest must be SHA-256")
        current = self.store.list_steps(task_id)[step_id]
        if current.state == StepState.COMPLETED and current.result_digest == result_digest:
            return
        self.store.transition_step(
            task_id,
            step_id,
            (StepState.RUNNING,),
            StepState.COMPLETED,
            result_digest=result_digest,
        )

    def fail(
        self,
        task: TaskPlan,
        step_id: str,
        reason: str,
        *,
        retryable: bool,
    ) -> StepState:
        current = self.store.list_steps(task.task_id)[step_id]
        next_state = (
            StepState.PENDING
            if retryable and current.attempt < current.max_attempts
            else StepState.FAILED
        )
        self.store.transition_step(
            task.task_id,
            step_id,
            (StepState.CLAIMED, StepState.RUNNING),
            next_state,
            last_error=reason,
            clear_assignment=next_state == StepState.PENDING,
        )
        self._task_state(
            task.task_id, TaskState.BLOCKED if next_state == StepState.PENDING else TaskState.FAILED
        )
        return next_state

    def recover_expired(
        self, task: TaskPlan, *, now: datetime | None = None
    ) -> tuple[str, ...]:
        now = now or datetime.now(timezone.utc)
        recovered: list[str] = []
        for step_id, current in self.store.list_steps(task.task_id).items():
            if (
                current.state == StepState.RUNNING
                and current.lease_expires_at is not None
                and current.lease_expires_at <= now
            ):
                self.fail(task, step_id, "lease_expired", retryable=True)
                recovered.append(step_id)
        return tuple(recovered)

    def recoverable(self, task_id: str, step_id: str) -> None:
        self.store.transition_step(
            task_id,
            step_id,
            (StepState.BLOCKED,),
            StepState.PENDING,
            clear_assignment=True,
        )