govoplan-workflow-engine/src/govoplan_workflow_engine/backend/recovery.py

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import json
from typing import Mapping

from sqlalchemy.orm import Session, sessionmaker

from govoplan_core.core.automation import ActionDefinition
from govoplan_core.core.recovery import (
    RecoveryGuaranteeError,
    RecoveryMode,
    RecoveryOperation,
    RecoveryPlan,
    RecoveryStatus,
)
from govoplan_core.core.recovery_runtime import (
    DurableRecoveryOperation,
    RecoveryOperationBusy,
    RecoveryOperationStateConflict,
    begin_durable_recovery_operation,
    claim_durable_recovery_operation,
)
from govoplan_core.core.runtime_coordination import (
    LeaseClaim,
    acquire_lease,
    assert_lease_fence,
    process_runtime_identity,
    release_lease,
)
from govoplan_workflow_engine.backend.db.models import (
    WorkflowDefinitionRevision,
    WorkflowInstance,
    WorkflowInstanceStep,
)


class WorkflowRecoveryError(RuntimeError):
    pass


class WorkflowRecoveryBusy(WorkflowRecoveryError):
    pass


class WorkflowRecoveryConflict(WorkflowRecoveryError):
    def __init__(self, operation_id: str, status: str) -> None:
        self.operation_id = operation_id
        self.status = status
        super().__init__(
            f"Workflow action recovery is {status}; reconcile it first"
        )


@dataclass(frozen=True, slots=True)
class WorkflowRecoveryDeclaration:
    operation_type: str
    mode: RecoveryMode | None
    boundaries: tuple[Mapping[str, object], ...]
    verification: tuple[str, ...]


WORKFLOW_RECOVERY_OPERATIONS = (
    WorkflowRecoveryDeclaration(
        operation_type="instance.state-transition",
        mode=RecoveryMode.ATOMIC,
        boundaries=(
            {
                "name": "worker-claim",
                "classification": "fenced",
                "resume": "a current process identity and distributed fence are required",
            },
            {
                "name": "instance-transition",
                "classification": "atomic",
                "resume": "the pinned revision, step, event, and instance state commit together",
            },
        ),
        verification=(
            "verify the pinned definition revision and current step",
            "verify the process still owns the distributed fence before commit",
        ),
    ),
    WorkflowRecoveryDeclaration(
        operation_type="activity.external-effect",
        mode=None,
        boundaries=(
            {
                "name": "action-preview",
                "classification": "recomputable",
                "resume": "revalidate capability, authority, input, and preview",
            },
            {
                "name": "provider-dispatch",
                "classification": "outcome-unknown",
                "resume": "verify by stable provider idempotency key before retry",
            },
            {
                "name": "workflow-projection",
                "classification": "verified",
                "resume": "commit the local projection with terminal recovery evidence",
            },
        ),
        verification=(
            "verify the pinned definition and canonical action request hashes",
            "verify the provider result and every announced effect",
            "reconcile an uncertain provider outcome before continuation",
        ),
    ),
)


def workflow_session_factory(session: Session) -> sessionmaker[Session]:
    bind = session.get_bind()
    if bind is None:
        raise WorkflowRecoveryError("Workflow recovery requires a database bind")
    return sessionmaker(bind=bind, expire_on_commit=False)


def canonical_sha256(value: object) -> str:
    encoded = json.dumps(
        value,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=True,
        default=str,
    ).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


@dataclass(slots=True)
class WorkflowActionRecovery:
    operation: DurableRecoveryOperation | None
    operation_id: str
    mode: RecoveryMode
    call_number: int
    replayed: bool

    def checkpoint_dispatch(
        self,
        session: Session,
        *,
        instance: WorkflowInstance,
        step: WorkflowInstanceStep,
        action_key: str,
        capability_name: str,
    ) -> None:
        if self.operation is None:
            return
        previous = dict(step.handoff or {})
        step.handoff = {
            **previous,
            "kind": "module_action",
            "state": "running",
            "action_key": action_key,
            "capability": capability_name,
            "allowed_actions": ["cancel"],
            "recovery": {
                "operation_id": self.operation_id,
                "mode": self.mode.value,
                "status": RecoveryStatus.RUNNING.value,
                "call_number": self.call_number,
                "boundary": "provider-dispatch",
                "requires_attention": False,
            },
        }
        step.status = "running"
        instance.status = "running"
        try:
            session.commit()
            self.operation.checkpoint(
                kind="provider-dispatch",
                summary="Provider dispatch crossed the recoverable effect boundary",
                evidence={
                    "effect_started": self.mode != RecoveryMode.ATOMIC,
                    "action_key": action_key,
                    "capability": capability_name,
                    "call_number": self.call_number,
                },
            )
        except Exception as exc:
            session.rollback()
            raise WorkflowRecoveryError(
                "Workflow action dispatch evidence could not be persisted"
            ) from exc

    def commit_conclusive_result(
        self,
        session: Session,
        *,
        provider_state: str,
        result_sha256: str,
        observed_effects_sha256: str,
    ) -> None:
        if self.operation is None:
            return
        evidence = {
            "verified": True,
            "checks": {
                "provider_state": provider_state,
                "result_sha256": result_sha256,
                "observed_effects_sha256": observed_effects_sha256,
                "call_number": self.call_number,
            },
        }
        if self.mode == RecoveryMode.ATOMIC:
            self.operation.commit_atomic_success(session, evidence=evidence)
        else:
            self.operation.commit_verified_success(session, evidence=evidence)

    def commit_unknown(
        self,
        session: Session,
        *,
        error_type: str,
        message: str,
    ) -> None:
        if self.operation is None:
            raise WorkflowRecoveryError(
                "A replayed action cannot acquire an unknown outcome"
            )
        try:
            session.commit()
            self.operation.unresolved(
                status=RecoveryStatus.OUTCOME_UNKNOWN,
                summary="The provider acknowledgement was not conclusive",
                evidence={
                    "effect_started": self.mode != RecoveryMode.ATOMIC,
                    "error_type": error_type,
                    "call_number": self.call_number,
                },
                failure_summary=message,
            )
        except Exception as exc:
            session.rollback()
            raise WorkflowRecoveryError(
                "Workflow action uncertainty could not be recorded"
            ) from exc

    def commit_definitive_failure(
        self,
        session: Session,
        *,
        summary: str,
        error_type: str,
    ) -> None:
        if self.operation is None:
            return
        evidence = {
            "verified": True,
            "checks": {
                "effect_started": False,
                "error_type": error_type,
                "call_number": self.call_number,
            },
        }
        if self.mode == RecoveryMode.ATOMIC:
            self.operation.commit_atomic_failure(
                session,
                summary=summary,
                evidence=evidence,
            )
        else:
            self.operation.fail(summary=summary, evidence=evidence)
            session.commit()


def begin_workflow_action_recovery(
    session: Session,
    *,
    instance: WorkflowInstance,
    step: WorkflowInstanceStep,
    revision: WorkflowDefinitionRevision,
    definition: ActionDefinition,
    capability_name: str,
    request_idempotency_key: str,
    action_input: Mapping[str, object],
    preview_payload: Mapping[str, object],
    backup_reference: str | None = None,
    approval_reference: str | None = None,
) -> WorkflowActionRecovery:
    previous_recovery = (
        step.handoff.get("recovery")
        if isinstance(step.handoff, Mapping)
        else None
    )
    call_number = 1
    if isinstance(previous_recovery, Mapping):
        call_number = max(
            1,
            int(
                previous_recovery.get("next_call_number")
                or previous_recovery.get("call_number")
                or 1
            ),
        )
    mode = RecoveryMode(definition.recovery_mode)
    plan = _action_recovery_plan(
        definition,
        mode=mode,
        backup_reference=backup_reference,
        approval_reference=approval_reference,
    )
    request_sha256 = canonical_sha256(dict(action_input))
    preview_sha256 = canonical_sha256(dict(preview_payload))
    access_context_sha256 = canonical_sha256(instance.authorization_)
    action_contract_sha256 = canonical_sha256(
        {
            "action_key": definition.action_key,
            "owner_module": definition.owner_module,
            "contract_version": definition.contract_version,
            "recovery_mode": definition.recovery_mode,
            "recovery_verification": list(definition.recovery_verification),
            "expected_effect_keys": list(definition.expected_effect_keys),
        }
    )
    session.commit()
    try:
        started = begin_durable_recovery_operation(
            workflow_session_factory(session),
            identity=process_runtime_identity(),
            module_id="workflow_engine",
            operation_type="activity.external-effect",
            idempotency_key=(
                f"workflow-step:{step.id}:action-call:{call_number}"
            ),
            request={
                "tenant_id": instance.tenant_id,
                "instance_id": instance.id,
                "step_id": step.id,
                "node_id": step.node_id,
                "definition_revision_id": revision.id,
                "definition_hash": revision.content_hash,
                "action_key": definition.action_key,
                "capability": capability_name,
                "action_contract_sha256": action_contract_sha256,
                "action_input_sha256": request_sha256,
                "preview_sha256": preview_sha256,
                "provider_idempotency_sha256": canonical_sha256(
                    request_idempotency_key
                ),
                "call_number": call_number,
            },
            recovery_plan=plan,
            precondition_evidence={
                "definition_hash": revision.content_hash,
                "action_contract_sha256": action_contract_sha256,
                "action_input_sha256": request_sha256,
                "preview_sha256": preview_sha256,
                "access_context_sha256": access_context_sha256,
                "call_number": call_number,
            },
            lease_resource_key=f"workflow:step:{step.id}",
            lease_ttl_seconds=300,
            resource_type="workflow_instance_step",
            resource_id=step.id,
            metadata={
                "resources": ["postgresql", "module-provider"],
                "workflow_instance_id": instance.id,
                "action_key": definition.action_key,
                "capability": capability_name,
            },
        )
    except RecoveryOperationBusy as exc:
        raise WorkflowRecoveryBusy(
            "Another runtime owns this Workflow action step"
        ) from exc
    except RecoveryOperationStateConflict as exc:
        raise WorkflowRecoveryConflict(exc.operation_id, exc.status) from exc
    except (RecoveryGuaranteeError, RuntimeError, ValueError) as exc:
        raise WorkflowRecoveryError(
            "The recovery ledger is unavailable; the module action did not start"
        ) from exc
    return WorkflowActionRecovery(
        operation=started.operation,
        operation_id=started.operation_id,
        mode=mode,
        call_number=call_number,
        replayed=started.replayed,
    )


def claim_workflow_action_recovery(
    session: Session,
    *,
    operation_id: str,
) -> DurableRecoveryOperation:
    try:
        return claim_durable_recovery_operation(
            workflow_session_factory(session),
            identity=process_runtime_identity(),
            operation_id=operation_id,
            lease_ttl_seconds=300,
        )
    except (RecoveryGuaranteeError, RuntimeError) as exc:
        raise WorkflowRecoveryError(
            "The Workflow action outcome cannot currently be reconciled"
        ) from exc


def reconcile_stale_workflow_action_recovery(
    session: Session,
    *,
    operation_id: str,
) -> str:
    try:
        handle = claim_durable_recovery_operation(
            workflow_session_factory(session),
            identity=process_runtime_identity(),
            operation_id=operation_id,
            lease_ttl_seconds=300,
        )
    except RecoveryOperationBusy as exc:
        raise WorkflowRecoveryBusy(
            "Another runtime still owns this Workflow action step"
        ) from exc
    except RecoveryOperationStateConflict as exc:
        return exc.status
    except (RecoveryGuaranteeError, RuntimeError) as exc:
        raise WorkflowRecoveryError(
            "The stale Workflow action fence could not be reconciled"
        ) from exc
    handle.release_unresolved()
    session.expire_all()
    operation = session.get(RecoveryOperation, operation_id)
    if operation is None:
        raise WorkflowRecoveryError(
            "The reconciled Workflow recovery operation is unavailable"
        )
    return str(operation.status)


def workflow_action_recovery_state(
    session: Session,
    *,
    operation_id: str,
) -> dict[str, object] | None:
    operation = session.get(RecoveryOperation, operation_id)
    if operation is None or operation.module_id != "workflow_engine":
        return None
    return {
        "operation_id": operation.id,
        "mode": operation.mode,
        "status": operation.status,
        "requires_attention": operation.status
        in {
            RecoveryStatus.OUTCOME_UNKNOWN.value,
            RecoveryStatus.RECOVERY_REQUIRED.value,
            RecoveryStatus.MANUAL_INTERVENTION.value,
        },
        "failure_summary": operation.failure_summary,
    }


def acquire_workflow_state_fence(
    session: Session,
    *,
    resource_key: str,
    ttl_seconds: int = 120,
) -> LeaseClaim | None:
    identity = process_runtime_identity()
    return acquire_lease(
        session,
        installation_id=identity.installation_id,
        resource_key=resource_key,
        holder_node_id=identity.node_id,
        holder_incarnation=identity.incarnation,
        ttl_seconds=ttl_seconds,
        metadata={"module_id": "workflow_engine", "kind": "state-transition"},
    )


def release_workflow_state_fence(session: Session, claim: LeaseClaim) -> None:
    assert_lease_fence(session, claim)
    release_lease(session, claim)


def _action_recovery_plan(
    definition: ActionDefinition,
    *,
    mode: RecoveryMode,
    backup_reference: str | None,
    approval_reference: str | None,
) -> RecoveryPlan:
    if mode == RecoveryMode.SNAPSHOT_RESTORE and not backup_reference:
        raise WorkflowRecoveryError(
            "Snapshot-restore actions require a pinned recovery backup reference"
        )
    if mode == RecoveryMode.IRREVERSIBLE and not approval_reference:
        raise WorkflowRecoveryError(
            "Irreversible actions require a pinned approval reference"
        )
    return RecoveryPlan(
        mode=mode,
        preconditions=(
            "the pinned Workflow revision and action contract hashes are present",
            "the current authority and provider capability were revalidated",
            "the action step owns a distributed execution fence",
        ),
        compensation_steps=(
            "invoke the provider-declared compensation action",
            "verify the domain invariant after compensation",
        )
        if mode == RecoveryMode.COMPENSATION
        else (),
        forward_recovery_steps=(
            "inspect the provider by stable idempotency key",
            "record whether the announced effect occurred",
            "continue or retry only after the outcome is proven",
        )
        if mode == RecoveryMode.FORWARD_RECOVERY
        else (),
        verification_steps=tuple(definition.recovery_verification),
        backup_reference=backup_reference,
        approval_reference=approval_reference,
    )


__all__ = [
    "WORKFLOW_RECOVERY_OPERATIONS",
    "WorkflowActionRecovery",
    "WorkflowRecoveryDeclaration",
    "WorkflowRecoveryBusy",
    "WorkflowRecoveryConflict",
    "WorkflowRecoveryError",
    "acquire_workflow_state_fence",
    "begin_workflow_action_recovery",
    "canonical_sha256",
    "claim_workflow_action_recovery",
    "release_workflow_state_fence",
    "reconcile_stale_workflow_action_recovery",
    "workflow_action_recovery_state",
    "workflow_session_factory",
]