gateway_contract.py

"""Public-safe protocol shared by the AutoML gateway and private backend."""

from __future__ import annotations

import json
import re
import stat
import zipfile
from pathlib import Path, PurePosixPath
from typing import Any, Mapping, Sequence

import pandas as pd

from gateway_routing import SUPPORTED_IMAGE_PREDICTORS


PROTOCOL_VERSION = "1.1"
MAX_UPLOAD_BYTES = 2 * 1024 * 1024 * 1024
MAX_PACKAGE_EXPANDED_BYTES = 10 * 1024 * 1024 * 1024
MAX_PACKAGE_MEMBERS = 20_000
MAX_MANIFEST_BYTES = 1024 * 1024

ERROR_CODES = frozenset(
    {
        "AUTH_REQUIRED",
        "PROTOCOL_MISMATCH",
        "INVALID_INPUT",
        "LIMIT_EXCEEDED",
        "BACKEND_WAKING",
        "BACKEND_TIMEOUT",
        "BACKEND_UNAVAILABLE",
        "JOB_REJECTED",
        "JOB_FAILED",
        "STORAGE_FAILED",
        "INTERNAL_ERROR",
    }
)

LIMITS = {
    "time_limit": (60, 3600),
    "trials": (1, 4),
    "image_epochs": (1, 100),
    "image_batch_size": (1, 128),
    "image_imgsz": (32, 1024),
}

_SENSITIVE_MESSAGE = re.compile(
    r"(?:Traceback|hf_[A-Za-z0-9]{6,}|/(?:data|tmp|home|Users)/|[A-Za-z]:\\)",
    re.IGNORECASE,
)
_GENERIC_INTERNAL_MESSAGE = "The backend could not complete the request."


class ContractError(ValueError):
    """Raised when an untrusted gateway/backend value violates the protocol."""


def _mapping(value: str | Mapping[str, Any]) -> dict[str, Any]:
    if isinstance(value, str):
        try:
            decoded = json.loads(value)
        except json.JSONDecodeError as exc:
            raise ContractError("Request must be valid JSON.") from exc
    else:
        decoded = value
    if not isinstance(decoded, Mapping):
        raise ContractError("Request must be a JSON object.")
    return dict(decoded)


def _bounded_int(payload: dict[str, Any], key: str, default: int) -> int:
    try:
        value = int(payload.get(key, default))
    except (TypeError, ValueError) as exc:
        raise ContractError(f"{key} must be an integer.") from exc
    lower, upper = LIMITS[key]
    if value < lower or value > upper:
        raise ContractError(f"{key} must be between {lower} and {upper}.")
    return value


def _required_text(payload: dict[str, Any], key: str) -> str:
    value = str(payload.get(key) or "").strip()
    if not value:
        raise ContractError(f"{key} is required.")
    if len(value) > 512:
        raise ContractError(f"{key} is too long.")
    payload[key] = value
    return value


def validate_request(kind: str, value: str | Mapping[str, Any]) -> dict[str, Any]:
    """Validate a versioned request before any backend model work starts."""

    payload = _mapping(value)
    if payload.get("protocol_version") != PROTOCOL_VERSION:
        raise ContractError(
            f"Gateway/backend protocol mismatch; expected {PROTOCOL_VERSION}."
        )
    if kind not in {"existing", "uploaded", "predict"}:
        raise ContractError("Unsupported request kind.")

    profile = _required_text(payload, "execution_profile")
    if profile not in {"cpu", "gpu"}:
        raise ContractError("execution_profile must be cpu or gpu.")
    mode = _required_text(payload, "routing_mode")
    if mode not in {"dual", "single-gpu"}:
        raise ContractError("routing_mode must be dual or single-gpu.")
    _required_text(payload, "route_reason_code")

    if kind == "existing":
        _required_text(payload, "dataset")
    elif kind == "uploaded":
        data_type = _required_text(payload, "data_type")
        if data_type not in {"Tabular CSV", "Image Dataset"}:
            raise ContractError("data_type must be Tabular CSV or Image Dataset.")

    if kind != "predict":
        for key, default in (
            ("time_limit", 300),
            ("trials", 1),
            ("image_epochs", 1),
            ("image_batch_size", 64),
            ("image_imgsz", 64),
        ):
            payload[key] = _bounded_int(payload, key, default)
        try:
            test_size = float(payload.get("test_size", 0.2))
        except (TypeError, ValueError) as exc:
            raise ContractError("test_size must be numeric.") from exc
        if not 0.1 <= test_size <= 0.5:
            raise ContractError("test_size must be between 0.1 and 0.5.")
        payload["test_size"] = test_size
    elif "image_batch_size" in payload:
        payload["image_batch_size"] = _bounded_int(
            payload, "image_batch_size", 16
        )
    return payload


def _safe_message(code: str, message: str) -> str:
    value = " ".join(str(message or "").split())[:500]
    if code == "INTERNAL_ERROR" or not value or _SENSITIVE_MESSAGE.search(value):
        return _GENERIC_INTERNAL_MESSAGE
    return value


def failure_envelope(
    code: str, message: str, retryable: bool = False
) -> dict[str, Any]:
    """Return a stable error result without secret or local-path disclosure."""

    normalized = code if code in ERROR_CODES else "INTERNAL_ERROR"
    return {
        "ok": False,
        "protocol_version": PROTOCOL_VERSION,
        "status": "failed",
        "error": {
            "code": normalized,
            "message": _safe_message(normalized, message),
            "retryable": bool(retryable),
        },
    }


def _frame(value: Any) -> pd.DataFrame:
    if isinstance(value, pd.DataFrame):
        return value.copy()
    if value is None:
        return pd.DataFrame()
    try:
        return pd.DataFrame(value)
    except Exception:
        return pd.DataFrame([{"result": str(value)}])


def success_envelope(outputs: Sequence[Any]) -> dict[str, Any]:
    """Normalize the existing seven training outputs into JSON-safe fields."""

    if len(outputs) != 7:
        raise ContractError("Backend training output must contain seven values.")
    status, metrics, predictions, summary, tools, reasoning, package = outputs
    frame = _frame(predictions)
    records = json.loads(frame.to_json(orient="records", date_format="iso"))
    return {
        "ok": True,
        "protocol_version": PROTOCOL_VERSION,
        "status": str(status or "complete"),
        "metrics_markdown": str(metrics or ""),
        "prediction_columns": [str(column) for column in frame.columns],
        "prediction_records": records,
        "summary": str(summary or ""),
        "tools_used": str(tools or ""),
        "reasoning": str(reasoning or ""),
        "package_available": bool(package),
    }


def prediction_envelope(outputs: Sequence[Any]) -> dict[str, Any]:
    """Normalize the existing three prediction outputs into the shared schema."""

    if len(outputs) != 3:
        raise ContractError("Backend prediction output must contain three values.")
    status, summary, predictions = outputs
    frame = _frame(predictions)
    records = json.loads(frame.to_json(orient="records", date_format="iso"))
    return {
        "ok": not str(status).startswith("❌"),
        "protocol_version": PROTOCOL_VERSION,
        "status": str(status or "complete"),
        "metrics_markdown": "",
        "prediction_columns": [str(column) for column in frame.columns],
        "prediction_records": records,
        "summary": str(summary or ""),
        "tools_used": "",
        "reasoning": "",
        "package_available": False,
    }


def predictions_frame(envelope: Mapping[str, Any]) -> pd.DataFrame:
    """Recreate a result DataFrame using the declared stable column order."""

    columns = [str(value) for value in envelope.get("prediction_columns") or []]
    records = envelope.get("prediction_records") or []
    if not isinstance(records, list):
        raise ContractError("prediction_records must be a list.")
    return pd.DataFrame(records, columns=columns or None)


def check_upload_size(path: str | Path, max_bytes: int = MAX_UPLOAD_BYTES) -> int:
    """Validate one local upload without reading it into memory."""

    candidate = Path(path)
    if not candidate.is_file():
        raise ContractError("Uploaded file was not found.")
    size = int(candidate.stat().st_size)
    if size > int(max_bytes):
        raise ContractError("Each uploaded file must be 2 GiB or smaller.")
    return size


def _safe_zip_member(info: zipfile.ZipInfo) -> None:
    name = info.filename
    if not name or "\\" in name:
        raise ContractError("Model package contains an unsafe archive path.")
    path = PurePosixPath(name)
    if path.is_absolute() or ".." in path.parts:
        raise ContractError("Model package contains an unsafe archive path.")
    mode = (info.external_attr >> 16) & 0xFFFF
    if stat.S_ISLNK(mode):
        raise ContractError("Model package links are not allowed.")
    if info.flag_bits & 0x1:
        raise ContractError("Encrypted ZIP members are not allowed.")


def inspect_package_manifest(path: str | Path) -> tuple[dict[str, Any], str]:
    """Read exactly one package manifest without extracting executable state."""

    package = Path(path)
    check_upload_size(package)
    try:
        with zipfile.ZipFile(package) as archive:
            members = archive.infolist()
            if len(members) > MAX_PACKAGE_MEMBERS:
                raise ContractError("Model package contains too many files.")
            expanded = 0
            manifests: list[zipfile.ZipInfo] = []
            for info in members:
                _safe_zip_member(info)
                expanded += int(info.file_size)
                if expanded > MAX_PACKAGE_EXPANDED_BYTES:
                    raise ContractError("Model package expands beyond the allowed size.")
                if PurePosixPath(info.filename).name == "automl_model_manifest.json":
                    manifests.append(info)
            if len(manifests) != 1:
                raise ContractError(
                    "Model package must contain exactly one automl_model_manifest.json."
                )
            manifest_info = manifests[0]
            if manifest_info.file_size > MAX_MANIFEST_BYTES:
                raise ContractError("Model package manifest is too large.")
            try:
                manifest = json.loads(archive.read(manifest_info).decode("utf-8"))
            except (UnicodeDecodeError, json.JSONDecodeError) as exc:
                raise ContractError("Model package manifest must be valid JSON.") from exc
    except zipfile.BadZipFile as exc:
        raise ContractError("Model package must be a valid ZIP file.") from exc
    if not isinstance(manifest, dict):
        raise ContractError("Model package manifest must be a JSON object.")
    artifact_type = str(manifest.get("artifact_type") or "").strip().lower()
    if not artifact_type:
        predictor = str(manifest.get("predictor") or "").strip()
        artifact_type = (
            "image" if predictor in SUPPORTED_IMAGE_PREDICTORS else "tabular"
        )
    if artifact_type not in {"tabular", "image"}:
        raise ContractError(
            f"Unsupported artifact_type in model manifest: {artifact_type!r}."
        )
    predictor = str(manifest.get("predictor") or "").strip()
    if (
        artifact_type == "image"
        and predictor
        and predictor not in SUPPORTED_IMAGE_PREDICTORS
    ):
        raise ContractError(
            f"Unsupported image predictor in model manifest: {predictor!r}."
        )
    return manifest, artifact_type


def load_gateway_catalog(path: str | Path) -> dict[str, Any]:
    """Load the metadata-only catalog shipped with the public gateway."""

    try:
        catalog = json.loads(Path(path).read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise ContractError("Gateway dataset catalog could not be loaded.") from exc
    if not isinstance(catalog, dict) or not isinstance(catalog.get("datasets"), dict):
        raise ContractError("Gateway catalog datasets must be a JSON object.")
    if int(catalog.get("schema_version", 1)) != 1:
        raise ContractError("Unsupported gateway catalog schema version.")
    normalized: dict[str, Any] = {"schema_version": 1, "datasets": {}}
    for name, metadata in catalog["datasets"].items():
        if not isinstance(metadata, dict):
            raise ContractError("Gateway catalog dataset metadata must be an object.")
        targets = metadata.get("targets") or []
        if not isinstance(targets, list) or not all(isinstance(item, str) for item in targets):
            raise ContractError("Gateway catalog targets must be a string list.")
        normalized["datasets"][str(name)] = {
            "targets": targets,
            "image": bool(metadata.get("image", False)),
        }
    return normalized


__all__ = [
    "PROTOCOL_VERSION",
    "MAX_UPLOAD_BYTES",
    "ERROR_CODES",
    "ContractError",
    "validate_request",
    "failure_envelope",
    "success_envelope",
    "prediction_envelope",
    "predictions_frame",
    "check_upload_size",
    "inspect_package_manifest",
    "load_gateway_catalog",
]