DataOps-platform/app/core/mcp/gateway.py

"""Machine policy and compound scheduling operations for MCP."""

import json
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime

from app.core.orchestration.compilers.kestra import compile_kestra_flow
from app.core.orchestration.spec import (
    validate_schedule_plan,
    validate_workflow_spec,
    workflow_spec_hash,
)

ALLOWED_SCHEDULING_TOOLS = {
    "create_candidate_plan",
    "deploy_disabled_version",
    "run_canary",
    "promote_candidate",
    "pause_schedule",
    "retry_failed_execution",
    "backfill_bounded_window",
    "rollback_to_previous_version",
}


@dataclass(frozen=True)
class GatewayLimits:
    max_concurrency: int = 10
    max_timeout_seconds: int = 7200
    max_retry_attempts: int = 3
    max_backfill_days: int = 7
    max_backfill_runs: int = 20


class SchedulingPolicy:
    _EDITOR_TOOLS = {
        "create_candidate_plan",
        "deploy_disabled_version",
        "run_canary",
    }
    _SCHEDULER_TOOLS = ALLOWED_SCHEDULING_TOOLS

    def __init__(self, limits=None):
        self.limits = limits or GatewayLimits()

    def authorize_tool(
        self,
        identity,
        tool_name,
        *,
        domain,
        environment,
    ):
        if tool_name not in ALLOWED_SCHEDULING_TOOLS:
            raise PermissionError("tool is not exposed by scheduling policy")
        identity.require_domain(domain)
        identity.require_environment(environment)
        if identity.has_any_role({"admin", "scheduler"}):
            allowed = self._SCHEDULER_TOOLS
        elif identity.has_any_role({"editor"}):
            allowed = self._EDITOR_TOOLS
        else:
            allowed = set()
        if tool_name not in allowed:
            raise PermissionError("role is not allowed to use this tool")

    def validate_plan_limits(self, plan):
        if plan.get("max_concurrency", 0) > self.limits.max_concurrency:
            raise ValueError("max concurrency exceeds gateway limit")
        if plan.get("timeout_seconds", 0) > self.limits.max_timeout_seconds:
            raise ValueError("timeout exceeds gateway limit")
        retry = plan.get("retry") or {}
        if retry.get("max_attempts", 0) > self.limits.max_retry_attempts:
            raise ValueError("retry attempts exceed gateway limit")
        backfill = plan.get("backfill") or {}
        if backfill.get("max_days", 0) > self.limits.max_backfill_days:
            raise ValueError("backfill days exceed gateway limit")
        if backfill.get("max_runs", 0) > self.limits.max_backfill_runs:
            raise ValueError("backfill runs exceed gateway limit")


class SchedulingGateway:
    def __init__(
        self,
        *,
        plans,
        audit,
        policy=None,
        engine=None,
        token_issuer=None,
        compiler=None,
    ):
        self.plans = plans
        self.audit = audit
        self.policy = policy or SchedulingPolicy()
        self.engine = engine
        self.token_issuer = token_issuer
        self.compiler = compiler or compile_kestra_flow

    @staticmethod
    def _event(identity, action, decision, domain, environment, **detail):
        return {
            "subject": identity.subject,
            "roles": sorted(identity.roles),
            "business_domain": domain,
            "environment": environment,
            "correlation_id": identity.correlation_id,
            "action": action,
            "decision": decision,
            "detail": detail,
        }

    @staticmethod
    def _idempotency_key(value, label):
        if not isinstance(value, str) or not value.strip() or len(value) > 500:
            raise ValueError(f"{label} idempotency key is required")
        return value.strip()

    @staticmethod
    def _timestamp(value, label):
        if not isinstance(value, str) or not value.strip():
            raise ValueError(f"{label} is required")
        source = value.strip()
        if source.endswith("Z"):
            source = f"{source[:-1]}+00:00"
        try:
            parsed = datetime.fromisoformat(source)
        except ValueError as exc:
            raise ValueError(f"{label} must be an ISO-8601 timestamp") from exc
        if parsed.tzinfo is None:
            raise ValueError(f"{label} must include a timezone")
        return parsed.astimezone(UTC)

    def _require_engine(self):
        if self.engine is None:
            raise RuntimeError("scheduling engine is not configured")
        return self.engine

    def _candidate(self, candidate_id, business_domain, environment):
        candidate = self.plans.get_candidate(candidate_id)
        if (
            candidate.get("business_domain") != business_domain
            or candidate.get("environment") != environment
        ):
            raise PermissionError("candidate is outside identity scope")
        return candidate

    def _deployment(self, candidate_id, business_domain, environment):
        self._candidate(candidate_id, business_domain, environment)
        deployment = self.plans.get_deployment(candidate_id)
        if not isinstance(deployment, dict):
            raise ValueError("candidate has no deployment")
        if not deployment.get("namespace") or not deployment.get("flow_id"):
            raise ValueError("candidate deployment is incomplete")
        return deployment

    def _record(
        self,
        identity,
        action,
        decision,
        business_domain,
        environment,
        **detail,
    ):
        self.audit.record(
            self._event(
                identity,
                action,
                decision,
                business_domain,
                environment,
                **detail,
            )
        )

    def create_candidate_plan(
        self,
        identity,
        *,
        business_domain,
        environment,
        workflow_spec,
        schedule_plan,
        context_text=None,
    ):
        del context_text  # External metadata is untrusted and never policy input.
        action = "create_candidate_plan"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        self.policy.validate_plan_limits(schedule_plan)
        spec = validate_workflow_spec(workflow_spec)
        plan = validate_schedule_plan(schedule_plan)
        result = self.plans.create_candidate(
            {
                "dataflow_uid": spec["dataflow_uid"],
                "business_domain": business_domain,
                "environment": environment,
                "workflow_spec": spec,
                "schedule_plan": plan,
                "workflow_hash": workflow_spec_hash(spec),
                "status": "candidate",
                "created_by": identity.subject,
                "correlation_id": identity.correlation_id,
            }
        )
        self.audit.record(
            self._event(
                identity,
                action,
                "allowed",
                business_domain,
                environment,
                candidate_id=result["candidate_id"],
                workflow_hash=result["workflow_hash"],
            )
        )
        return result

    def deploy_disabled_version(
        self,
        identity,
        *,
        candidate_id,
        business_domain,
        environment,
    ):
        action = "deploy_disabled_version"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        candidate = self._candidate(candidate_id, business_domain, environment)
        compiled = self.compiler(
            candidate["workflow_spec"],
            candidate["schedule_plan"],
            environment,
            int(candidate["version_no"]),
        )
        response = self._require_engine().deploy_disabled(compiled.yaml)
        deployment = {
            "candidate_id": candidate_id,
            "namespace": compiled.namespace,
            "flow_id": compiled.flow_id,
            "definition_hash": compiled.definition_hash,
            "engine_revision": (
                response.get("revision") if isinstance(response, dict) else None
            ),
            "status": "deployed_disabled",
        }
        result = self.plans.record_deployment(candidate_id, deployment)
        self._record(
            identity,
            action,
            "allowed",
            business_domain,
            environment,
            candidate_id=candidate_id,
            flow_id=compiled.flow_id,
            definition_hash=compiled.definition_hash,
        )
        return result

    def run_canary(
        self,
        identity,
        *,
        candidate_id,
        business_domain,
        environment,
        inputs,
    ):
        action = "run_canary"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        if not isinstance(inputs, dict):
            raise ValueError("canary inputs must be an object")
        if "dataops_task_tokens" in inputs:
            raise ValueError("reserved task token input cannot be supplied")
        if len(inputs) > 100:
            raise ValueError("canary inputs exceed gateway limit")
        try:
            input_size = len(
                json.dumps(
                    inputs,
                    sort_keys=True,
                    separators=(",", ":"),
                    ensure_ascii=False,
                ).encode("utf-8")
            )
        except (TypeError, ValueError) as exc:
            raise ValueError("canary inputs must be JSON serializable") from exc
        if input_size > 65536:
            raise ValueError("canary inputs exceed gateway size limit")

        candidate = self._candidate(candidate_id, business_domain, environment)
        unknown_inputs = sorted(
            set(inputs) - set(candidate["workflow_spec"]["parameters"])
        )
        if unknown_inputs:
            raise ValueError(
                f"canary inputs contain unknown parameters: {', '.join(unknown_inputs)}"
            )
        deployment = self._deployment(candidate_id, business_domain, environment)
        if self.token_issuer is None:
            raise RuntimeError("runner task token issuer is not configured")
        task_tokens = {}
        for node in candidate["workflow_spec"]["nodes"]:
            task_tokens[node["id"]] = self.token_issuer.issue(
                task_uid=str(uuid.uuid4()),
                dataflow_uid=candidate["dataflow_uid"],
                workflow_version=int(candidate["version_no"]),
                correlation_id=identity.correlation_id,
                node=node,
                write_authorized=bool(candidate.get("trusted_write_authorized", False)),
            )
        engine = self._require_engine()
        # Kestra rejects manual execution while a flow is disabled. Compiled
        # schedule triggers remain individually disabled, so the gateway opens
        # only a bounded manual-execution window and always closes it again.
        engine.activate(deployment["namespace"], deployment["flow_id"])
        try:
            response = engine.execute(
                deployment["namespace"],
                deployment["flow_id"],
                inputs={**inputs, "dataops_task_tokens": task_tokens},
            )
        finally:
            engine.deactivate(deployment["namespace"], deployment["flow_id"])
        execution_id = response.get("id") if isinstance(response, dict) else None
        if not execution_id:
            raise RuntimeError("canary execution did not return an id")
        evidence = {
            "candidate_id": candidate_id,
            "execution_id": execution_id,
            "status": "started",
            "sample_runs": 0,
        }
        result = self.plans.record_canary_execution(candidate_id, evidence)
        self._record(
            identity,
            action,
            "allowed",
            business_domain,
            environment,
            candidate_id=candidate_id,
            execution_id=execution_id,
        )
        return result

    def promote_candidate(
        self,
        identity,
        *,
        candidate_id,
        business_domain,
        environment,
        idempotency_key,
    ):
        action = "promote_candidate"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        self._candidate(candidate_id, business_domain, environment)
        deployment = self._deployment(candidate_id, business_domain, environment)
        key = self._idempotency_key(idempotency_key, "promotion")
        canary = self.plans.get_canary_evidence(candidate_id)
        if (
            not isinstance(canary, dict)
            or canary.get("status") != "passed"
            or int(canary.get("sample_runs", 0)) < 1
            or not canary.get("verified_by")
        ):
            raise ValueError("trusted successful canary evidence is required")
        claimed, result = self.plans.claim_promotion(
            key,
            candidate_id,
            actor_subject=identity.subject,
            correlation_id=identity.correlation_id,
        )
        decision = "allowed" if claimed else "idempotent_replay"
        if claimed:
            active_lookup = getattr(self.plans, "get_active_deployment", None)
            previous = (
                active_lookup(candidate_id) if active_lookup is not None else None
            )
            try:
                self._require_engine().activate(
                    deployment["namespace"], deployment["flow_id"]
                )
                if previous is not None:
                    self._require_engine().deactivate(
                        previous["namespace"], previous["flow_id"]
                    )
                complete = getattr(self.plans, "complete_promotion", None)
                if complete is not None:
                    complete(key, candidate_id, result)
            except Exception:
                try:
                    self._require_engine().deactivate(
                        deployment["namespace"], deployment["flow_id"]
                    )
                    if previous is not None:
                        self._require_engine().activate(
                            previous["namespace"], previous["flow_id"]
                        )
                except Exception:
                    pass
                release = getattr(self.plans, "release_promotion", None)
                if release is not None:
                    release(key, candidate_id)
                raise
        self._record(
            identity,
            action,
            decision,
            business_domain,
            environment,
            candidate_id=candidate_id,
            idempotency_key=key,
        )
        return result

    def pause_schedule(
        self,
        identity,
        *,
        candidate_id,
        business_domain,
        environment,
    ):
        action = "pause_schedule"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        deployment = self._deployment(candidate_id, business_domain, environment)
        self._require_engine().deactivate(
            deployment["namespace"], deployment["flow_id"]
        )
        self.plans.mark_paused(candidate_id)
        result = {
            "candidate_id": candidate_id,
            "status": "paused",
            "flow_id": deployment["flow_id"],
        }
        self._record(
            identity,
            action,
            "allowed",
            business_domain,
            environment,
            candidate_id=candidate_id,
            flow_id=deployment["flow_id"],
        )
        return result

    def retry_failed_execution(
        self,
        identity,
        *,
        execution_id,
        business_domain,
        environment,
        max_retries=1,
    ):
        action = "retry_failed_execution"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        if (
            isinstance(max_retries, bool)
            or not isinstance(max_retries, int)
            or max_retries < 1
            or max_retries > self.policy.limits.max_retry_attempts
        ):
            raise ValueError("max retries exceed gateway limit")
        record = self.plans.get_execution_record(execution_id)
        if (
            record.get("business_domain") != business_domain
            or record.get("environment") != environment
        ):
            raise PermissionError("execution is outside identity scope")
        if str(record.get("status", "")).upper() != "FAILED":
            raise ValueError("only failed executions can be retried")
        if int(record.get("retry_count", 0)) >= max_retries:
            raise ValueError("execution retry limit is exhausted")
        result = self._require_engine().replay(execution_id)
        self.plans.record_retry(execution_id, result)
        self._record(
            identity,
            action,
            "allowed",
            business_domain,
            environment,
            execution_id=execution_id,
            replay_execution_id=(
                result.get("id") if isinstance(result, dict) else None
            ),
        )
        return result

    def backfill_bounded_window(
        self,
        identity,
        *,
        candidate_id,
        business_domain,
        environment,
        trigger_id,
        start,
        end,
    ):
        action = "backfill_bounded_window"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        candidate = self._candidate(candidate_id, business_domain, environment)
        deployment = self._deployment(candidate_id, business_domain, environment)
        start_at = self._timestamp(start, "backfill start")
        end_at = self._timestamp(end, "backfill end")
        if end_at <= start_at:
            raise ValueError("backfill end must be after start")
        window_days = (end_at - start_at).total_seconds() / 86400
        allowed_days = min(
            self.policy.limits.max_backfill_days,
            int(candidate["schedule_plan"]["backfill"]["max_days"]),
        )
        if window_days > allowed_days:
            raise ValueError("backfill window exceeds gateway limit")
        estimated_runs = int(
            self.plans.estimate_backfill_runs(candidate_id, start, end)
        )
        allowed_runs = min(
            self.policy.limits.max_backfill_runs,
            int(candidate["schedule_plan"]["backfill"]["max_runs"]),
        )
        if estimated_runs < 1 or estimated_runs > allowed_runs:
            raise ValueError("backfill run count exceeds gateway limit")
        response = self._require_engine().backfill(
            deployment["namespace"],
            deployment["flow_id"],
            trigger_id,
            start,
            end,
        )
        result = {
            **(response if isinstance(response, dict) else {}),
            "candidate_id": candidate_id,
            "estimated_runs": estimated_runs,
        }
        self.plans.record_backfill(candidate_id, result)
        self._record(
            identity,
            action,
            "allowed",
            business_domain,
            environment,
            candidate_id=candidate_id,
            trigger_id=trigger_id,
            estimated_runs=estimated_runs,
        )
        return result

    def rollback_to_previous_version(
        self,
        identity,
        *,
        candidate_id,
        business_domain,
        environment,
        idempotency_key,
    ):
        action = "rollback_to_previous_version"
        self.policy.authorize_tool(
            identity,
            action,
            domain=business_domain,
            environment=environment,
        )
        current = self._deployment(candidate_id, business_domain, environment)
        previous = self.plans.get_previous_deployment(candidate_id)
        if (
            not isinstance(previous, dict)
            or not previous.get("namespace")
            or not previous.get("flow_id")
        ):
            raise ValueError("previous deployment is not available")
        key = self._idempotency_key(idempotency_key, "rollback")
        claimed, result = self.plans.claim_rollback(
            key,
            candidate_id,
            actor_subject=identity.subject,
            correlation_id=identity.correlation_id,
        )
        if not claimed:
            self._record(
                identity,
                action,
                "idempotent_replay",
                business_domain,
                environment,
                candidate_id=candidate_id,
                idempotency_key=key,
            )
            return result

        engine = self._require_engine()
        engine.deactivate(current["namespace"], current["flow_id"])
        try:
            engine.activate(previous["namespace"], previous["flow_id"])
        except Exception:
            engine.activate(current["namespace"], current["flow_id"])
            release = getattr(self.plans, "release_rollback", None)
            if release is not None:
                release(key, candidate_id)
            raise
        complete = getattr(self.plans, "complete_rollback", None)
        if complete is not None:
            complete(key, candidate_id, result)
        self._record(
            identity,
            action,
            "allowed",
            business_domain,
            environment,
            candidate_id=candidate_id,
            active_candidate_id=previous.get("candidate_id"),
            idempotency_key=key,
        )
        return result