ThothII/harness/tht/jobs/models.py

"""Immutable, secret-free records for preprocessing execution."""

from __future__ import annotations

import re
from datetime import UTC, datetime
from pathlib import Path
from typing import Literal, Self

from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator


_JOB_KEY = re.compile(r"^[a-z][a-z0-9_-]{0,63}$")
_RUN_ID = re.compile(r"^[0-9a-f]{32}$")
_FINGERPRINT = re.compile(r"^sha256:[0-9a-f]{64}$")
JobStatus = Literal["pending", "running", "succeeded", "failed"]
StageStatus = Literal["pending", "running", "succeeded", "failed"]


def utc_now() -> datetime:
    return datetime.now(UTC)


def _validate_job_key(value: str) -> str:
    if not _JOB_KEY.fullmatch(value):
        raise ValueError("job identifiers must be lowercase filesystem-safe keys")
    return value


def _validate_run_id(value: str | None) -> str | None:
    if value is not None and not _RUN_ID.fullmatch(value):
        raise ValueError("run id must contain 32 lowercase hexadecimal characters")
    return value


class _FrozenModel(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid", validate_default=True)

    def model_copy(self, *, update=None, deep: bool = False) -> Self:
        data = self.model_dump(round_trip=True)
        if update:
            data.update(update)
        return type(self).model_validate(data)

    @field_serializer("*", when_used="json", check_fields=False)
    def serialize_utc(self, value):
        if isinstance(value, datetime):
            return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
        return value


class JobSpec(_FrozenModel):
    """Execution input. The local root is deliberately excluded from serialization."""

    workspace_id: str
    job_type: str
    workspace_root: Path = Field(exclude=True)
    spec_version: str = Field(min_length=1, max_length=64)
    pipeline_version: str = Field(min_length=1, max_length=64)
    config_fingerprint: str
    input_fingerprint: str
    stage_ids: tuple[str, ...]
    dry_run: bool = False
    resume_run_id: str | None = None

    _workspace_key = field_validator("workspace_id")(_validate_job_key)
    _job_type_key = field_validator("job_type")(_validate_job_key)
    _version_keys = field_validator("spec_version", "pipeline_version")(_validate_job_key)
    _resume_id = field_validator("resume_run_id")(_validate_run_id)
    _config_fingerprint = field_validator("config_fingerprint")(
        lambda value: value if _FINGERPRINT.fullmatch(value) else _invalid_fingerprint()
    )
    _input_fingerprint = field_validator("input_fingerprint")(
        lambda value: value if _FINGERPRINT.fullmatch(value) else _invalid_fingerprint()
    )
    _stage_ids = field_validator("stage_ids")(
        lambda values: tuple(_validate_job_key(value) for value in values)
    )

    def model_copy(self, *, update=None, deep: bool = False) -> Self:
        data = {
            "workspace_id": self.workspace_id,
            "job_type": self.job_type,
            "workspace_root": self.workspace_root,
            "spec_version": self.spec_version,
            "pipeline_version": self.pipeline_version,
            "config_fingerprint": self.config_fingerprint,
            "input_fingerprint": self.input_fingerprint,
            "stage_ids": self.stage_ids,
            "dry_run": self.dry_run,
            "resume_run_id": self.resume_run_id,
        }
        if update:
            data.update(update)
        return type(self).model_validate(data)

    def with_resume(self, run_id: str) -> "JobSpec":
        return self.model_copy(update={"resume_run_id": run_id})


class StageError(_FrozenModel):
    category: Literal["internal"] = "internal"
    code: Literal["stage_exception"] = "stage_exception"
    message: Literal["stage execution failed"] = "stage execution failed"


class StageRun(_FrozenModel):
    name: str
    status: StageStatus = "pending"
    started_at: datetime | None = None
    finished_at: datetime | None = None
    error: StageError | None = None

    _name_key = field_validator("name")(_validate_job_key)


class JobRun(_FrozenModel):
    """Durable checkpoint, persisted after every state transition."""

    schema_version: Literal[1] = 1
    run_id: str
    compatibility_fingerprint: str
    workspace_fingerprint: str
    job_type: str
    spec_version: str
    pipeline_version: str
    config_fingerprint: str
    input_fingerprint: str
    dry_run: bool
    status: JobStatus
    started_at: datetime
    finished_at: datetime | None = None
    resumed_from: str | None = None
    stages: tuple[StageRun, ...] = ()

    _run_id = field_validator("run_id")(_validate_run_id)
    _compatibility = field_validator("compatibility_fingerprint", "workspace_fingerprint")(
        lambda value: value if _FINGERPRINT.fullmatch(value) else _invalid_fingerprint()
    )
    _input_fingerprints = field_validator("config_fingerprint", "input_fingerprint")(
        lambda value: value if _FINGERPRINT.fullmatch(value) else _invalid_fingerprint()
    )
    _job_type = field_validator("job_type")(_validate_job_key)
    _persisted_versions = field_validator("spec_version", "pipeline_version")(_validate_job_key)
    _resumed_from = field_validator("resumed_from")(_validate_run_id)


class JobReport(JobRun):
    """Public machine-readable terminal report (contains no paths or stage outputs)."""


def _invalid_fingerprint():
    raise ValueError("fingerprint must be sha256 followed by 64 lowercase hexadecimal characters")