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

"""Bounded read-only business context for scheduling agents."""

from app.core.orchestration.spec import (
    validate_schedule_plan as normalize_schedule_plan,
)
from app.core.orchestration.spec import (
    validate_workflow_spec as normalize_workflow_spec,
)
from app.core.orchestration.spec import workflow_spec_hash


def _untrusted(value):
    if value is None:
        return None
    return {"trust": "untrusted", "value": str(value)[:2000]}


def _page(items, maximum):
    bounded = list(items)[:maximum]
    return {
        "items": bounded,
        "count": len(bounded),
        "truncated": len(items) > maximum,
    }


class ContextService:
    def __init__(self, repository):
        self.repository = repository

    def list_dataflows(self, identity, *, limit=50):
        if limit < 1 or limit > 50:
            raise ValueError("limit must be between 1 and 50")
        allowed = []
        for row in self.repository.list_dataflows():
            domain = row.get("business_domain")
            if (
                "*" not in identity.business_domains
                and domain not in identity.business_domains
            ):
                continue
            allowed.append(
                {
                    "uid": row.get("uid"),
                    "name": row.get("name"),
                    "business_domain": domain,
                    "description": _untrusted(row.get("description")),
                }
            )
        return _page(allowed, limit)

    def describe_dataflow(self, identity, dataflow_uid, *, business_domain):
        identity.require_domain(business_domain)
        row = self.repository.describe_dataflow(dataflow_uid)
        if row.get("business_domain") != business_domain:
            raise PermissionError("dataflow is outside identity scope")
        return {
            "uid": row.get("uid"),
            "name": row.get("name"),
            "business_domain": row.get("business_domain"),
            "status": row.get("status"),
            "description": _untrusted(row.get("description")),
        }

    def get_dataflow_dependencies(self, identity, dataflow_uid, *, business_domain):
        identity.require_domain(business_domain)
        safe = []
        for row in self.repository.get_dataflow_dependencies(dataflow_uid):
            if row.get("business_domain") != business_domain:
                continue
            safe.append(
                {
                    "uid": row.get("uid"),
                    "name": row.get("name"),
                    "relation": row.get("relation"),
                    "business_domain": row.get("business_domain"),
                }
            )
        return _page(safe, 100)

    def get_data_lineage(
        self,
        identity,
        dataflow_uid,
        *,
        business_domain,
        depth=3,
    ):
        identity.require_domain(business_domain)
        if (
            isinstance(depth, bool)
            or not isinstance(depth, int)
            or depth < 1
            or depth > 5
        ):
            raise ValueError("lineage depth must be between 1 and 5")
        safe = []
        for row in self.repository.get_data_lineage(dataflow_uid, depth):
            safe.append(
                {
                    "from_uid": row.get("from_uid"),
                    "to_uid": row.get("to_uid"),
                    "relation": row.get("relation"),
                    "description": _untrusted(row.get("description")),
                }
            )
        return _page(safe, 100)

    def list_datasource_capabilities(self, identity, business_domain):
        identity.require_domain(business_domain)
        safe = []
        for row in self.repository.list_datasource_capabilities(business_domain):
            safe.append(
                {
                    "uid": row.get("uid"),
                    "type": row.get("type"),
                    "purposes": list(row.get("purposes") or [])[:10],
                    "health": row.get("health"),
                    "capacity": {
                        key: value
                        for key, value in dict(row.get("capacity") or {}).items()
                        if key
                        in {
                            "available",
                            "pool_size",
                            "checked_out",
                            "max_concurrency",
                        }
                    },
                }
            )
        return _page(safe, 50)

    def get_datasource_pool_health(
        self,
        identity,
        data_source_uid,
        *,
        business_domain,
    ):
        identity.require_domain(business_domain)
        row = self.repository.get_datasource_pool_health(data_source_uid)
        return {
            "uid": row.get("uid"),
            "health": row.get("health"),
            "capacity": {
                key: value
                for key, value in dict(row.get("capacity") or {}).items()
                if key
                in {
                    "available",
                    "pool_size",
                    "checked_out",
                    "max_concurrency",
                }
            },
        }

    def get_execution_history(
        self,
        identity,
        dataflow_uid,
        *,
        business_domain,
        limit=20,
        days=7,
    ):
        identity.require_domain(business_domain)
        if limit < 1 or limit > 100:
            raise ValueError("limit must be between 1 and 100")
        if days < 1 or days > 30:
            raise ValueError("days must be between 1 and 30")
        safe = []
        for row in self.repository.get_execution_history(
            dataflow_uid,
            limit,
            days,
        ):
            safe.append(
                {
                    "status": row.get("status"),
                    "correlation_id": row.get("correlation_id"),
                    "started_at": row.get("started_at"),
                    "finished_at": row.get("finished_at"),
                    "safe_error": _untrusted(row.get("safe_error")),
                }
            )
        return _page(safe, limit)

    def get_sla_constraints(self, identity, dataflow_uid, *, business_domain):
        identity.require_domain(business_domain)
        row = self.repository.get_sla_constraints(dataflow_uid)
        return {
            key: row.get(key)
            for key in (
                "timezone",
                "max_duration_seconds",
                "deadline",
                "priority",
            )
            if key in row
        }

    def estimate_schedule_capacity(
        self,
        identity,
        *,
        business_domain,
        schedule_plan,
    ):
        identity.require_domain(business_domain)
        plan = normalize_schedule_plan(schedule_plan)
        row = self.repository.estimate_schedule_capacity(business_domain, plan)
        return {
            "feasible": bool(row.get("feasible", False)),
            "required_slots": row.get("required_slots"),
            "available_slots": row.get("available_slots"),
            "warnings": [
                _untrusted(value) for value in list(row.get("warnings") or [])[:20]
            ],
        }

    def validate_workflow_spec(self, identity, workflow_spec):
        del identity
        normalized = normalize_workflow_spec(workflow_spec)
        return {
            "valid": True,
            "workflow_spec": normalized,
            "workflow_hash": workflow_spec_hash(normalized),
        }

    def validate_schedule_plan(self, identity, schedule_plan):
        del identity
        return {
            "valid": True,
            "schedule_plan": normalize_schedule_plan(schedule_plan),
        }

    def simulate_schedule(
        self,
        identity,
        *,
        business_domain,
        schedule_plan,
    ):
        identity.require_domain(business_domain)
        plan = normalize_schedule_plan(schedule_plan)
        row = self.repository.simulate_schedule(business_domain, plan)
        return {
            "next_runs": list(row.get("next_runs") or [])[:20],
            "warnings": [
                _untrusted(value) for value in list(row.get("warnings") or [])[:20]
            ],
        }

    def compare_execution_results(
        self,
        identity,
        *,
        business_domain,
        baseline_execution_id,
        candidate_execution_id,
    ):
        identity.require_domain(business_domain)
        row = self.repository.compare_execution_results(
            baseline_execution_id,
            candidate_execution_id,
        )
        metrics = {
            key: value
            for key, value in dict(row.get("metrics") or {}).items()
            if key
            in {
                "row_count_delta",
                "duration_delta_seconds",
                "error_count_delta",
                "output_hash_match",
            }
            and (value is None or isinstance(value, (bool, int, float)))
        }
        return {
            "equivalent": bool(row.get("equivalent", False)),
            "metrics": metrics,
            "differences": [
                _untrusted(value) for value in list(row.get("differences") or [])[:50]
            ],
        }