"""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