DataOps-platform/tests/agent/test_planner_scenarios.py

import copy

import pytest

from app.core.mcp.identity import AgentIdentity
from app.core.orchestration.agent.evaluator import replay_candidate_plan
from app.core.orchestration.agent.planner import (
    OpenAICompatiblePlanningModel,
    SchedulingPlanner,
)

DATAFLOW_UID = "01900000-0000-7000-8000-000000000054"
SOURCE_UID = "01900000-0000-7000-8000-000000000055"


def identity():
    return AgentIdentity(
        subject="v54-planner",
        roles=frozenset({"scheduler"}),
        business_domains=frozenset({"sales"}),
        environments=frozenset({"test"}),
        correlation_id="01900000-0000-7000-8000-000000000056",
    )


def planning_request():
    return {
        "business_goal": "Refresh the customer summary before the 08:00 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-1",
    }


def model_plan():
    return {
        "schema_version": "1.0",
        "workflow_spec": {
            "schema_version": "1.0",
            "dataflow_uid": DATAFLOW_UID,
            "name": "Customer summary read-only canary",
            "nodes": [
                {
                    "id": "read_customers",
                    "type": "sql.query",
                    "data_source_uid": SOURCE_UID,
                    "purpose": "read",
                    "config": {
                        "statement": (
                            "SELECT customer_name FROM 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": 600,
            "retry": {"max_attempts": 1, "delay_seconds": 0},
            "backfill": {"max_days": 1, "max_runs": 1},
        },
        "decision_summary": "Use one bounded read-only canary.",
    }


class Model:
    provider = "test-provider"
    model_name = "test-model"

    def __init__(self, response=None):
        self.response = copy.deepcopy(response or model_plan())
        self.calls = []

    def generate(self, *, messages, response_schema, timeout_seconds):
        self.calls.append(
            {
                "messages": messages,
                "response_schema": response_schema,
                "timeout_seconds": timeout_seconds,
            }
        )
        return copy.deepcopy(self.response)


class FailingModel(Model):
    def __init__(self, error):
        super().__init__()
        self.error = error

    def generate(self, *, messages, response_schema, timeout_seconds):
        super().generate(
            messages=messages,
            response_schema=response_schema,
            timeout_seconds=timeout_seconds,
        )
        raise self.error


class ContextTools:
    def __init__(self):
        self.calls = []

    def call_tool(self, name, arguments):
        self.calls.append((name, arguments))
        responses = {
            "get_sla_constraints": {
                "timezone": "Asia/Shanghai",
                "deadline": "08:00",
                "max_duration_seconds": 1800,
            },
            "list_datasource_capabilities": {
                "items": [
                    {
                        "uid": SOURCE_UID,
                        "purposes": ["read"],
                        "health": "healthy",
                        "capacity": {"available": 4, "max_concurrency": 4},
                    }
                ],
                "count": 1,
                "truncated": False,
            },
            "get_execution_history": {
                "items": [{"status": "success"}],
                "count": 1,
                "truncated": False,
            },
            "validate_workflow_spec": {"valid": True},
            "validate_schedule_plan": {"valid": True},
            "estimate_schedule_capacity": {
                "feasible": True,
                "required_slots": 1,
                "available_slots": 4,
                "warnings": [],
            },
            "simulate_schedule": {
                "next_runs": [],
                "warnings": [],
            },
        }
        return copy.deepcopy(responses[name])


class SchedulingTools:
    def __init__(self):
        self.calls = []

    def call_tool(self, name, arguments):
        self.calls.append((name, copy.deepcopy(arguments)))
        responses = {
            "create_candidate_plan": {
                "candidate_id": "candidate-v54",
                "workflow_hash": "a" * 64,
                "status": "candidate",
            },
            "deploy_disabled_version": {
                "candidate_id": "candidate-v54",
                "status": "deployed_disabled",
            },
            "run_canary": {
                "candidate_id": "candidate-v54",
                "execution_id": "kestra-canary-v54",
                "status": "started",
            },
        }
        return copy.deepcopy(responses[name])


class FailingSchedulingTools(SchedulingTools):
    def __init__(self, failed_tool):
        super().__init__()
        self.failed_tool = failed_tool

    def call_tool(self, name, arguments):
        if name == self.failed_tool:
            self.calls.append((name, copy.deepcopy(arguments)))
            raise RuntimeError("upstream secret-looking text must not be audited")
        return super().call_tool(name, arguments)


class Verifier:
    def __init__(self):
        self.calls = []

    def verify(self, candidate_id, *, baseline_execution_id=None):
        self.calls.append((candidate_id, baseline_execution_id))
        return {
            "candidate_id": candidate_id,
            "execution_id": "kestra-canary-v54",
            "status": "passed",
            "sample_runs": 1,
            "verified_by": "v54-verifier",
            "verification_summary": {"equivalent": True},
        }


class Audit:
    def __init__(self):
        self.events = []

    def record(self, event):
        self.events.append(copy.deepcopy(event))


def test_planner_runs_fixed_read_only_canary_sequence_and_records_decision():
    model = Model()
    context = ContextTools()
    scheduling = SchedulingTools()
    verifier = Verifier()
    audit = Audit()
    planner = SchedulingPlanner(
        model=model,
        context_tools=context,
        scheduling_tools=scheduling,
        canary_verifier=verifier,
        audit=audit,
        model_timeout_seconds=15,
    )

    result = planner.run(identity(), planning_request())

    assert result["status"] == "canary_passed"
    assert result["candidate_id"] == "candidate-v54"
    assert result["promotion_recommendation"] == "shadow_only"
    assert result["formal_engine_changed"] is False
    assert [name for name, _ in context.calls] == [
        "get_sla_constraints",
        "list_datasource_capabilities",
        "get_execution_history",
        "validate_workflow_spec",
        "validate_schedule_plan",
        "estimate_schedule_capacity",
        "simulate_schedule",
    ]
    assert [name for name, _ in scheduling.calls] == [
        "create_candidate_plan",
        "deploy_disabled_version",
        "run_canary",
    ]
    assert verifier.calls == [("candidate-v54", "n8n-baseline-1")]
    assert model.calls[0]["timeout_seconds"] == 15
    assert model.calls[0]["response_schema"]["additionalProperties"] is False
    event = audit.events[-1]
    assert event["action"] == "ai_planning_closed_loop"
    assert event["decision"] == "allowed"
    assert event["model_provider"] == "test-provider"
    assert event["model_name"] == "test-model"
    assert event["prompt_version"]
    assert event["schema_version"] == "v54"
    assert len(event["context_hash"]) == 64
    assert event["detail"]["candidate_plan"]["schema_version"] == "1.0"
    assert event["detail"]["validation"]["capacity"]["feasible"] is True
    assert [item["tool"] for item in event["detail"]["action_results"]] == [
        "create_candidate_plan",
        "deploy_disabled_version",
        "run_canary",
        "verify_canary",
    ]


def test_model_timeout_degrades_to_fixed_schedule_without_mutating_tools():
    scheduling = SchedulingTools()
    audit = Audit()
    planner = SchedulingPlanner(
        model=FailingModel(TimeoutError("provider timeout")),
        context_tools=ContextTools(),
        scheduling_tools=scheduling,
        canary_verifier=Verifier(),
        audit=audit,
    )

    result = planner.run(identity(), planning_request())

    assert result == {
        "status": "degraded",
        "mode": "fixed_schedule",
        "reason_code": "model_unavailable",
        "action": "continue_published_schedule",
        "formal_engine_changed": False,
        "correlation_id": identity().correlation_id,
    }
    assert scheduling.calls == []
    assert audit.events[-1]["decision"] == "failed"
    assert audit.events[-1]["detail"]["reason_code"] == "model_unavailable"
    assert "provider timeout" not in str(audit.events[-1])


def test_invalid_or_write_model_plan_is_safely_rejected_before_scheduling():
    write_plan = model_plan()
    write_plan["workflow_spec"]["nodes"][0].update(
        {
            "type": "sql.execute",
            "purpose": "write",
            "idempotency": {
                "strategy": "deduplication_key",
                "key": "customer-summary-v54",
            },
        }
    )
    scheduling = SchedulingTools()
    audit = Audit()
    planner = SchedulingPlanner(
        model=Model(write_plan),
        context_tools=ContextTools(),
        scheduling_tools=scheduling,
        canary_verifier=Verifier(),
        audit=audit,
    )

    result = planner.run(identity(), planning_request())

    assert result["status"] == "rejected"
    assert result["reason_code"] == "invalid_model_output"
    assert result["action"] == "continue_published_schedule"
    assert scheduling.calls == []
    assert audit.events[-1]["decision"] == "rejected"
    assert "candidate_plan" not in audit.events[-1]["detail"]
    assert len(audit.events[-1]["detail"]["candidate_plan_hash"]) == 64


def test_target_conflict_is_rejected_without_calling_model_or_tools():
    request = planning_request()
    request["target_conflicts"] = [
        "two workflows claim the same customer_summary target"
    ]
    model = Model()
    context = ContextTools()
    scheduling = SchedulingTools()
    audit = Audit()
    planner = SchedulingPlanner(
        model=model,
        context_tools=context,
        scheduling_tools=scheduling,
        canary_verifier=Verifier(),
        audit=audit,
    )

    result = planner.run(identity(), request)

    assert result["status"] == "rejected"
    assert result["reason_code"] == "target_conflict"
    assert model.calls == []
    assert context.calls == []
    assert scheduling.calls == []
    assert audit.events[-1]["decision"] == "rejected"


def test_tool_failure_stops_at_disabled_candidate_and_audits_safe_result():
    scheduling = FailingSchedulingTools("deploy_disabled_version")
    audit = Audit()
    planner = SchedulingPlanner(
        model=Model(),
        context_tools=ContextTools(),
        scheduling_tools=scheduling,
        canary_verifier=Verifier(),
        audit=audit,
    )

    result = planner.run(identity(), planning_request())

    assert result["status"] == "degraded"
    assert result["reason_code"] == "tool_failure"
    assert result["candidate_id"] == "candidate-v54"
    assert result["action"] == "continue_published_schedule"
    assert [name for name, _ in scheduling.calls] == [
        "create_candidate_plan",
        "deploy_disabled_version",
    ]
    assert audit.events[-1]["decision"] == "failed"
    assert audit.events[-1]["detail"]["action_results"][0]["tool"] == (
        "create_candidate_plan"
    )
    assert "secret-looking" not in str(audit.events[-1])


def test_openai_compatible_model_requests_deterministic_json_and_supplies_schema():
    class Message:
        content = '{"schema_version":"1.0"}'

    class Choice:
        message = Message()

    class Completions:
        def __init__(self):
            self.calls = []

        def create(self, **kwargs):
            self.calls.append(kwargs)
            return type("Response", (), {"choices": [Choice()]})()

    completions = Completions()
    client = type(
        "Client",
        (),
        {"chat": type("Chat", (), {"completions": completions})()},
    )()
    model = OpenAICompatiblePlanningModel(
        client=client,
        provider="deepseek",
        model_name="deepseek-chat",
    )

    result = model.generate(
        messages=[{"role": "user", "content": "plan"}],
        response_schema={"type": "object", "additionalProperties": False},
        timeout_seconds=12,
    )

    assert result == '{"schema_version":"1.0"}'
    call = completions.calls[0]
    assert call["model"] == "deepseek-chat"
    assert call["temperature"] == 0
    assert call["timeout"] == 12
    assert call["response_format"] == {"type": "json_object"}
    assert "OUTPUT_JSON_SCHEMA" in call["messages"][-1]["content"]
    assert '"additionalProperties":false' in call["messages"][-1]["content"]


def test_saved_candidate_can_be_replayed_offline_without_model_or_tools():
    first = replay_candidate_plan(model_plan(), planning_request())
    second = replay_candidate_plan(model_plan(), planning_request())

    assert first == second
    assert first["allowed"] is True
    assert first["schema_version"] == "v54"
    assert len(first["context_hash"]) == 64
    assert len(first["candidate_hash"]) == 64
    assert len(first["workflow_hash"]) == 64


def test_planner_monitors_canary_until_trusted_terminal_evidence():
    class ProgressiveVerifier(Verifier):
        def __init__(self):
            super().__init__()
            self.responses = [
                {
                    "candidate_id": "candidate-v54",
                    "status": "started",
                    "sample_runs": 0,
                },
                {
                    "candidate_id": "candidate-v54",
                    "status": "started",
                    "sample_runs": 0,
                },
                {
                    "candidate_id": "candidate-v54",
                    "status": "passed",
                    "sample_runs": 1,
                    "verified_by": "v54-verifier",
                },
            ]

        def verify(self, candidate_id, *, baseline_execution_id=None):
            self.calls.append((candidate_id, baseline_execution_id))
            return copy.deepcopy(self.responses.pop(0))

    verifier = ProgressiveVerifier()
    sleeps = []
    ticks = iter([0.0, 0.1, 0.2, 0.3])
    planner = SchedulingPlanner(
        model=Model(),
        context_tools=ContextTools(),
        scheduling_tools=SchedulingTools(),
        canary_verifier=verifier,
        audit=Audit(),
        canary_timeout_seconds=5,
        canary_poll_interval_seconds=0.25,
        sleep=sleeps.append,
        monotonic=lambda: next(ticks),
    )

    result = planner.run(identity(), planning_request())

    assert result["status"] == "canary_passed"
    assert len(verifier.calls) == 3
    assert sleeps == [0.25, 0.25]


@pytest.mark.parametrize(
    "mutate",
    [
        lambda candidate: candidate["schedule_plan"]["retry"].update(
            {"max_attempts": 1000000}
        ),
        lambda candidate: candidate["workflow_spec"]["nodes"][0]["config"].update(
            {"password": "model-requested-secret"}
        ),
        lambda candidate: candidate.update(
            {"backfill_partitions": ["already-successful-2026-07-18"]}
        ),
    ],
)
def test_fixed_scenarios_reject_unbounded_secret_or_model_backfill_advice(
    mutate,
):
    candidate = model_plan()
    mutate(candidate)
    scheduling = SchedulingTools()
    planner = SchedulingPlanner(
        model=Model(candidate),
        context_tools=ContextTools(),
        scheduling_tools=scheduling,
        canary_verifier=Verifier(),
        audit=Audit(),
    )

    result = planner.run(identity(), planning_request())

    assert result["status"] == "rejected"
    assert result["reason_code"] == "invalid_model_output"
    assert scheduling.calls == []