cdsl-cad/backend/app/cad_agent/domain/state.py

"""Finite workflow state machine.  This module has no persistence imports."""

from __future__ import annotations

from dataclasses import dataclass, replace
from enum import StrEnum

from .errors import ErrorCode, WorkflowError


class TaskPhase(StrEnum):
    DRAFTING_REQUIREMENTS_DOCUMENT = "DRAFTING_REQUIREMENTS_DOCUMENT"
    DRAFTING_COMPLETION_TARGET = "DRAFTING_COMPLETION_TARGET"
    COMPILING_REQUIREMENTS = "COMPILING_REQUIREMENTS"
    DRAFTING_MODELING_PLAN = "DRAFTING_MODELING_PLAN"
    AWAITING_ACTION = "AWAITING_ACTION"
    ACTION_PENDING = "ACTION_PENDING"
    CANDIDATE_BUILDING = "CANDIDATE_BUILDING"
    CANDIDATE_REVIEW = "CANDIDATE_REVIEW"
    FINAL_VALIDATION = "FINAL_VALIDATION"
    WAITING_RETRY = "WAITING_RETRY"
    WAITING_FOR_USER = "WAITING_FOR_USER"
    CANCELLED = "CANCELLED"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"


@dataclass(frozen=True, slots=True)
class PendingAction:
    action_id: str
    working_head: str
    intent: str
    requirement_ids: tuple[str, ...]
    atomic_id: str
    expected_change: str
    contract_hash: str
    idempotency_key: str


@dataclass(frozen=True, slots=True)
class TaskState:
    task_id: str
    phase: TaskPhase
    version: int
    active_revision: str = ""
    pending_action: PendingAction | None = None
    candidate_id: str = ""
    candidate_stage_id: str = ""
    repair_required: bool = False
    last_error: ErrorCode | None = None
    retry_from_phase: TaskPhase | None = None
    requirements_spec_path: str = ""
    requirements_document_path: str = ""
    completion_target_path: str = ""
    modeling_plan_path: str = ""
    clarification_path: str = ""
    requirements_contract_path: str = ""

    @property
    def working_head(self) -> str:
        return f"{self.task_id}:{self.active_revision or 'root'}:v{self.version}"


# Legal state transitions.  Events are intentionally terse persistence-neutral
# names used by command handlers and architecture tests.
_TRANSITIONS: dict[tuple[TaskPhase, str], TaskPhase] = {
    (TaskPhase.DRAFTING_REQUIREMENTS_DOCUMENT, "image_observed"): TaskPhase.DRAFTING_REQUIREMENTS_DOCUMENT,
    (TaskPhase.DRAFTING_REQUIREMENTS_DOCUMENT, "requirements_document_written"): TaskPhase.DRAFTING_COMPLETION_TARGET,
    (TaskPhase.DRAFTING_COMPLETION_TARGET, "completion_target_written"): TaskPhase.COMPILING_REQUIREMENTS,
    (TaskPhase.COMPILING_REQUIREMENTS, "requirements_compiled"): TaskPhase.DRAFTING_MODELING_PLAN,
    (TaskPhase.DRAFTING_MODELING_PLAN, "modeling_plan_written"): TaskPhase.AWAITING_ACTION,
    (TaskPhase.DRAFTING_REQUIREMENTS_DOCUMENT, "waiting_for_user"): TaskPhase.WAITING_FOR_USER,
    # User clarifications are durable task evidence. Resume on the same task
    # so its frozen request remains authoritative
    # instead of turning a clarification into a new CAD request.
    (TaskPhase.WAITING_FOR_USER, "requirements_clarified"): TaskPhase.DRAFTING_REQUIREMENTS_DOCUMENT,
    (TaskPhase.AWAITING_ACTION, "action_proposed"): TaskPhase.ACTION_PENDING,
    (TaskPhase.AWAITING_ACTION, "diagnosis_recorded"): TaskPhase.AWAITING_ACTION,
    (TaskPhase.AWAITING_ACTION, "rollback"): TaskPhase.AWAITING_ACTION,
    (TaskPhase.ACTION_PENDING, "candidate_started"): TaskPhase.CANDIDATE_BUILDING,
    # Runtime semantic preflight happens before candidate staging. It must
    # leave the checkpoint untouched and return control to the author for a
    # fresh observation/action decision, never masquerade as a format error.
    (TaskPhase.ACTION_PENDING, "runtime_precondition_rejected"): TaskPhase.AWAITING_ACTION,
    (TaskPhase.ACTION_PENDING, "diagnosis_recorded"): TaskPhase.ACTION_PENDING,
    (TaskPhase.ACTION_PENDING, "diagnosis_return_to_action_selection"): TaskPhase.AWAITING_ACTION,
    (TaskPhase.CANDIDATE_BUILDING, "candidate_built"): TaskPhase.CANDIDATE_REVIEW,
    (TaskPhase.CANDIDATE_BUILDING, "candidate_rejected"): TaskPhase.ACTION_PENDING,
    (TaskPhase.CANDIDATE_REVIEW, "candidate_accepted"): TaskPhase.AWAITING_ACTION,
    (TaskPhase.CANDIDATE_REVIEW, "candidate_rejected"): TaskPhase.AWAITING_ACTION,
    (TaskPhase.AWAITING_ACTION, "final_requested"): TaskPhase.FINAL_VALIDATION,
    (TaskPhase.FINAL_VALIDATION, "final_accepted"): TaskPhase.COMPLETED,
    (TaskPhase.FINAL_VALIDATION, "final_repair"): TaskPhase.AWAITING_ACTION,
}
_TRANSITIONS.update({
    (phase, "best_effort_completed"): TaskPhase.COMPLETED
    for phase in TaskPhase
    if phase not in {TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED}
})
_TRANSITIONS.update({
    (phase, "failed"): TaskPhase.FAILED
    for phase in TaskPhase
    if phase not in {TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED}
})
_TRANSITIONS.update({
    (phase, "cancelled"): TaskPhase.CANCELLED
    for phase in TaskPhase
    if phase not in {TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED}
})
_RETRY_RESUMABLE_PHASES = frozenset({
    TaskPhase.DRAFTING_REQUIREMENTS_DOCUMENT,
    TaskPhase.DRAFTING_COMPLETION_TARGET,
    TaskPhase.COMPILING_REQUIREMENTS,
    TaskPhase.DRAFTING_MODELING_PLAN,
    TaskPhase.AWAITING_ACTION,
    TaskPhase.ACTION_PENDING,
    TaskPhase.CANDIDATE_BUILDING,
    TaskPhase.CANDIDATE_REVIEW,
    TaskPhase.FINAL_VALIDATION,
})
_TRANSITIONS.update({
    (TaskPhase.WAITING_RETRY, f"resume_{phase.value.lower()}"): phase
    for phase in _RETRY_RESUMABLE_PHASES
})
_TRANSITIONS.update({
    (phase, "waiting_retry"): TaskPhase.WAITING_RETRY
    for phase in TaskPhase
    if phase not in {TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED, TaskPhase.WAITING_RETRY}
})
def legal_transitions() -> dict[tuple[TaskPhase, str], TaskPhase]:
    """Return a copy of the protocol transition table for architecture tests."""
    return dict(_TRANSITIONS)


def retry_resume_event(state: TaskState) -> str | None:
    """Return the only lossless resume event for a parked service failure."""
    if state.phase != TaskPhase.WAITING_RETRY or state.retry_from_phase not in _RETRY_RESUMABLE_PHASES:
        return None
    return f"resume_{state.retry_from_phase.value.lower()}"


def transition(state: TaskState, event: str, *, pending_action: PendingAction | None | object = ..., active_revision: str | None = None, candidate_id: str | None = None, candidate_stage_id: str | None = None, repair_required: bool | None = None, error: ErrorCode | None = None, requirements_spec_path: str | None = None, requirements_document_path: str | None = None, completion_target_path: str | None = None, modeling_plan_path: str | None = None, clarification_path: str | None = None, requirements_contract_path: str | None = None) -> TaskState:
    """Apply one legal transition and advance optimistic-concurrency version."""
    target = _TRANSITIONS.get((state.phase, event))
    if target is None:
        raise ValueError(f"Illegal v3 transition: {state.phase.value} --{event}--> ?")
    if state.phase == TaskPhase.WAITING_RETRY and event.startswith("resume_") and state.retry_from_phase != target:
        raise ValueError("WAITING_RETRY resume event does not match its persisted source phase")
    next_pending = state.pending_action if pending_action is ... else pending_action
    if target in {TaskPhase.AWAITING_ACTION, TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED}:
        next_pending = None
    return replace(
        state,
        phase=target,
        version=state.version + 1,
        active_revision=state.active_revision if active_revision is None else active_revision,
        pending_action=next_pending,
        candidate_id="" if target in {TaskPhase.AWAITING_ACTION, TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED} else state.candidate_id if candidate_id is None else candidate_id,
        candidate_stage_id="" if target in {TaskPhase.AWAITING_ACTION, TaskPhase.COMPLETED, TaskPhase.FAILED, TaskPhase.CANCELLED} else state.candidate_stage_id if candidate_stage_id is None else candidate_stage_id,
        repair_required=state.repair_required if repair_required is None else repair_required,
        last_error=error,
        retry_from_phase=state.phase if target == TaskPhase.WAITING_RETRY else None,
        requirements_spec_path=state.requirements_spec_path if requirements_spec_path is None else requirements_spec_path,
        requirements_document_path=state.requirements_document_path if requirements_document_path is None else requirements_document_path,
        completion_target_path=state.completion_target_path if completion_target_path is None else completion_target_path,
        modeling_plan_path=state.modeling_plan_path if modeling_plan_path is None else modeling_plan_path,
        clarification_path=state.clarification_path if clarification_path is None else clarification_path,
        requirements_contract_path=state.requirements_contract_path if requirements_contract_path is None else requirements_contract_path,
    )


def reject_stale_head(state: TaskState, supplied_head: str) -> WorkflowError | None:
    if supplied_head != state.working_head:
        return WorkflowError(
            ErrorCode.STALE_WORKING_HEAD,
            "The command was bound to an obsolete working head.",
            details={"expected_working_head": state.working_head, "supplied_working_head": supplied_head},
        )
    return None