DataOps-platform/tests/integration/test_agent_kestra_canary.py

"""Opt-in V54 Agent -> DataOps MCP -> Kestra -> Runner Canary acceptance."""

import os

import pytest
import requests
from neo4j import GraphDatabase
from sqlalchemy import create_engine, text

from app.core.common.identifiers import new_governance_uid
from app.core.data_source.credentials import (
    CredentialCodec,
    DataSourceCredentialRepository,
)
from app.core.data_source.definitions import DataSourceDefinitionRepository
from app.core.data_source.models import (
    DataSourceCredential,
    DataSourceDefinition,
)
from app.core.mcp.canary import CanaryVerifier
from app.core.mcp.context import ContextService
from app.core.mcp.gateway import SchedulingGateway
from app.core.mcp.identity import AgentIdentity
from app.core.mcp.persistence import (
    PostgresAuditSink,
    PostgresSchedulingPlanStore,
)
from app.core.mcp.servers import create_context_mcp, create_scheduling_mcp
from app.core.orchestration.agent.planner import SchedulingPlanner
from app.core.orchestration.engines.kestra import KestraAdapter
from app.runner.auth import TaskTokenIssuer

pytestmark = pytest.mark.integration
DATABASE_URL = "postgresql://dataops:dataops-test-password@127.0.0.1:15432/dataops"
KESTRA_URL = "http://127.0.0.1:18080/api/v1"
KESTRA_AUTH = ("admin@dataops.local", "DataOpsKestra1!")


class RegisteredMcp:
    def __init__(self, name, **_kwargs):
        self.name = name
        self.tools = {}

    def tool(self):
        def register(function):
            self.tools[function.__name__] = function
            return function

        return register


class ToolClient:
    def __init__(self, server):
        self.server = server

    def call_tool(self, name, arguments):
        return self.server.tools[name](**arguments)


class ContextRepository:
    def __init__(self, source_uid):
        self.source_uid = source_uid

    def get_sla_constraints(self, _dataflow_uid):
        return {
            "timezone": "Asia/Shanghai",
            "deadline": "08:00",
            "max_duration_seconds": 120,
            "priority": "low",
        }

    def list_datasource_capabilities(self, _business_domain):
        return [
            {
                "uid": self.source_uid,
                "type": "postgresql",
                "purposes": ["read"],
                "health": "healthy",
                "capacity": {"available": 2, "max_concurrency": 2},
            }
        ]

    def get_execution_history(self, _dataflow_uid, _limit, _days):
        return [{"status": "success", "correlation_id": new_governance_uid()}]

    def estimate_schedule_capacity(self, _business_domain, _schedule_plan):
        return {
            "feasible": True,
            "required_slots": 1,
            "available_slots": 2,
            "warnings": [],
        }

    def simulate_schedule(self, _business_domain, _schedule_plan):
        return {
            "next_runs": [],
            "warnings": ["manual trigger has no deterministic next run"],
        }


class DeterministicModel:
    provider = "offline-fixture"
    model_name = "v54-read-only-scenario"

    def __init__(self, dataflow_uid, source_uid):
        self.dataflow_uid = dataflow_uid
        self.source_uid = source_uid

    def generate(self, **_kwargs):
        return {
            "schema_version": "1.0",
            "workflow_spec": {
                "schema_version": "1.0",
                "dataflow_uid": self.dataflow_uid,
                "name": "V54 read-only agent Canary",
                "nodes": [
                    {
                        "id": "read_customers",
                        "type": "sql.query",
                        "data_source_uid": self.source_uid,
                        "purpose": "read",
                        "config": {
                            "statement": (
                                "SELECT customer_name FROM acceptance_customers "
                                "WHERE id >= :minimum_id ORDER BY id"
                            ),
                            "parameters": {"minimum_id": "${parameters.minimum_id}"},
                        },
                    }
                ],
                "edges": [],
                "parameters": {"minimum_id": {"type": "integer", "required": True}},
            },
            "schedule_plan": {
                "schema_version": "1.0",
                "timezone": "Asia/Shanghai",
                "triggers": [{"type": "manual"}],
                "max_concurrency": 1,
                "conflict_policy": "skip",
                "timeout_seconds": 120,
                "retry": {"max_attempts": 1, "delay_seconds": 0},
                "backfill": {"max_days": 1, "max_runs": 1},
            },
            "decision_summary": "Run one read-only low-frequency Canary.",
        }


class BaselineComparator:
    def compare(self, baseline_execution_id, candidate_execution_id):
        assert baseline_execution_id == "n8n-baseline-v54"
        assert candidate_execution_id
        return {
            "equivalent": True,
            "metrics": {
                "row_count_delta": 0,
                "output_hash_match": True,
            },
        }


def test_agent_executes_one_read_only_kestra_canary_without_promotion():
    if os.environ.get("RUN_V54_AGENT_L3") != "1":
        pytest.skip("set RUN_V54_AGENT_L3=1 for the V54 Agent acceptance")

    engine = create_engine(DATABASE_URL, pool_pre_ping=True)
    graph = GraphDatabase.driver(
        "bolt://127.0.0.1:17687",
        auth=("neo4j", "Passw0rd"),
    )
    source_uid = new_governance_uid()
    dataflow_uid = new_governance_uid()
    correlation_id = new_governance_uid()
    candidate_id = None
    deployment = None
    codec = CredentialCodec.from_base64(
        "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
        "v1",
    )
    identity = AgentIdentity(
        subject="v54-planner-l3",
        roles=frozenset({"scheduler"}),
        business_domains=frozenset({"sales"}),
        environments=frozenset({"test"}),
        correlation_id=correlation_id,
    )
    store = PostgresSchedulingPlanStore(engine)
    kestra = KestraAdapter(
        base_url=KESTRA_URL,
        username=KESTRA_AUTH[0],
        password=KESTRA_AUTH[1],
    )

    try:
        with engine.begin() as connection:
            sealed = DataSourceCredentialRepository(codec).create_version(
                connection,
                data_source_uid=source_uid,
                credential=DataSourceCredential(
                    "source_reader",
                    "source-test-password",
                ),
                actor_uid=None,
            )
        with graph.session() as session:
            DataSourceDefinitionRepository(session).save(
                DataSourceDefinition(
                    uid=source_uid,
                    name_en="v54-agent-canary-source",
                    database_type="postgresql",
                    host="source-postgres",
                    port=5432,
                    database="acceptance",
                    credential_ref=source_uid,
                    credential_version=sealed.credential_version,
                    pool_size=1,
                    max_overflow=1,
                )
            )

        context_server = create_context_mcp(
            ContextService(ContextRepository(source_uid)),
            identity=identity,
            fastmcp_cls=RegisteredMcp,
        )
        scheduling_server = create_scheduling_mcp(
            SchedulingGateway(
                plans=store,
                audit=PostgresAuditSink(engine),
                engine=kestra,
                token_issuer=TaskTokenIssuer(
                    "dataops-local-runner-task-token-secret-change-me",
                    ttl_seconds=120,
                ),
            ),
            identity=identity,
            fastmcp_cls=RegisteredMcp,
        )
        planner = SchedulingPlanner(
            model=DeterministicModel(dataflow_uid, source_uid),
            context_tools=ToolClient(context_server),
            scheduling_tools=ToolClient(scheduling_server),
            canary_verifier=CanaryVerifier(
                plans=store,
                engine=kestra,
                comparator=BaselineComparator(),
                verifier_id="v54-agent-canary-verifier",
            ),
            audit=PostgresAuditSink(engine),
            canary_timeout_seconds=30,
            canary_poll_interval_seconds=0.25,
        )

        result = planner.run(
            identity,
            {
                "business_goal": "Validate the customer read path before SLA",
                "dataflow_uid": dataflow_uid,
                "business_domain": "sales",
                "environment": "test",
                "available_window": {
                    "start": "2026-07-20T06:00:00+08:00",
                    "end": "2026-07-20T08:00:00+08:00",
                },
                "authorized_resources": [source_uid],
                "canary_inputs": {"minimum_id": 1},
                "baseline_execution_id": "n8n-baseline-v54",
            },
        )

        candidate_id = result["candidate_id"]
        deployment = store.get_deployment(candidate_id)
        assert result["status"] == "canary_passed"
        assert result["promotion_recommendation"] == "shadow_only"
        assert result["formal_engine_changed"] is False
        assert (
            requests.get(
                f"{KESTRA_URL}/main/flows/"
                f"{deployment['namespace']}/{deployment['flow_id']}",
                auth=KESTRA_AUTH,
                timeout=20,
            ).json()["disabled"]
            is True
        )

        with engine.connect() as connection:
            runner_successes = connection.execute(
                text(
                    "SELECT COUNT(*) FROM public.runner_task_executions "
                    "WHERE dataflow_uid = CAST(:uid AS uuid) "
                    "AND status = 'success'"
                ),
                {"uid": dataflow_uid},
            ).scalar_one()
            audit_row = (
                connection.execute(
                    text(
                        "SELECT model_provider, model_name, prompt_version, "
                        "schema_version, context_hash, decision, decision_detail "
                        "FROM public.workflow_plan_audits "
                        "WHERE action = 'ai_planning_closed_loop' "
                        "AND correlation_id = CAST(:correlation_id AS uuid)"
                    ),
                    {"correlation_id": correlation_id},
                )
                .mappings()
                .one()
            )
        assert runner_successes == 1
        assert audit_row["model_provider"] == "offline-fixture"
        assert audit_row["schema_version"] == "v54"
        assert audit_row["decision"] == "allowed"
        assert audit_row["decision_detail"]["result"]["formal_engine_changed"] is False
    finally:
        if deployment is not None:
            requests.delete(
                f"{KESTRA_URL}/main/flows/"
                f"{deployment['namespace']}/{deployment['flow_id']}",
                auth=KESTRA_AUTH,
                timeout=20,
            )
        with graph.session() as session:
            DataSourceDefinitionRepository(session).delete(source_uid)
        with engine.begin() as connection:
            if candidate_id is not None:
                connection.execute(
                    text(
                        "DELETE FROM public.workflow_plan_audits "
                        "WHERE workflow_version_id = CAST(:id AS uuid)"
                    ),
                    {"id": candidate_id},
                )
                connection.execute(
                    text(
                        "DELETE FROM public.dataflow_workflow_versions "
                        "WHERE id = CAST(:id AS uuid)"
                    ),
                    {"id": candidate_id},
                )
            connection.execute(
                text(
                    "DELETE FROM public.runner_task_executions "
                    "WHERE dataflow_uid = CAST(:uid AS uuid)"
                ),
                {"uid": dataflow_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.datasource_credential_audit_events "
                    "WHERE data_source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.datasource_credentials "
                    "WHERE data_source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
        graph.close()
        engine.dispose()