lambdagent/tests/test_pipeline.py

"""
Tests for engine.pipeline — 产物 Gate 逐级验收流水线 (P0+P1 骨架).

覆盖：
  1. 静态校验（forward-ref consumes、duplicate id）
  2. 受限 gate DSL 求值器（功能 + 安全：拒绝属性/import/裸名）
  3. resolve_artifacts 运行时契约校验（required 缺失）
  4. run_pipeline 端到端（fake executor）：通过 / 多上游 consumes / gate 失败 retry→halt / skip
  5. 每阶段 verdict + manifest + write-ahead state
  6. resume_pipeline 跳过已通过阶段
"""
import json
import os
import shutil
import tempfile

import pytest

from agentpaas.engine.pipeline import (
    load_pipeline, validate_pipeline, eval_rule_gate, resolve_artifacts,
    run_pipeline, resume_pipeline, _parse_judge_response,
    estimate_pipeline_cost, build_judge_term, term_judge, JudgeRequest,
    Stage, ArtifactSpec, ResolvedArtifact, GateError, PipelineError,
    StageStatus, PipelineStatus, StageExecResult, JudgeResult,
)

_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_DEMO_PRESET = os.path.join(_REPO, "agentexample", "pipelines", "research-demo.yml")


def _ollama_up() -> bool:
    import urllib.request
    try:
        urllib.request.urlopen("http://127.0.0.1:11434/api/tags", timeout=2)
        return True
    except Exception:
        return False


@pytest.fixture
def agent_dir():
    d = tempfile.mkdtemp(prefix="test_pipeline_")
    yield d
    shutil.rmtree(d, ignore_errors=True)


# ============================================================
# helpers — fake executors that write artifacts (no LLM)
# ============================================================

def _writer(files: dict):
    """Return a stage executor that writes {relpath: content} into the stage ws."""
    def _exec(ec):
        for rel, content in files.items():
            p = os.path.join(ec.stage_ws, rel)
            os.makedirs(os.path.dirname(p), exist_ok=True)
            with open(p, "w", encoding="utf-8") as f:
                f.write(content)
        return StageExecResult(output="ok", status="completed")
    return _exec


# ============================================================
# 1. Static validation
# ============================================================

class TestStaticValidation:

    def test_valid_pipeline(self):
        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "produces": [{"name": "x.md", "path": "final/x.md"}],
             "gate": {"kind": "rule", "rule": 'has("x.md")'}},
            {"id": "b", "consumes": ["a"], "gate": {"kind": "rule", "rule": "True"}},
        ]}}
        p = load_pipeline(cfg)
        assert validate_pipeline(p) == []

    def test_forward_consumes_rejected(self):
        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "consumes": ["b"]},   # b not yet declared
            {"id": "b"},
        ]}}
        errs = validate_pipeline(load_pipeline(cfg))
        assert any("forward" in e for e in errs)

    def test_duplicate_stage_id_rejected(self):
        cfg = {"pipeline": {"id": "p", "stages": [{"id": "a"}, {"id": "a"}]}}
        errs = validate_pipeline(load_pipeline(cfg))
        assert any("duplicate" in e for e in errs)

    def test_run_rejects_invalid(self, agent_dir):
        cfg = {"pipeline": {"id": "p", "stages": [{"id": "a", "consumes": ["z"]}]}}
        with pytest.raises(PipelineError):
            run_pipeline(cfg, "in", agent_dir, stage_executor=_writer({}))

    def test_unsafe_rule_caught_statically(self, agent_dir):
        """§15 P2-8: a banned gate expression fails validation BEFORE the executor runs."""
        ran = {"n": 0}

        def exec_spy(ec):
            ran["n"] += 1
            return StageExecResult(output="ok", status="completed")

        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "gate": {"kind": "rule", "rule": 'open("/etc/passwd")'}},
        ]}}
        errs = validate_pipeline(load_pipeline(cfg))
        assert any("not allowed" in e for e in errs)
        with pytest.raises(PipelineError):
            run_pipeline(cfg, "in", agent_dir, stage_executor=exec_spy)
        assert ran["n"] == 0   # executor never ran


# ============================================================
# 2. Bounded gate DSL evaluator
# ============================================================

class TestGateEvaluator:

    def _resolved(self, tmp, name="plan.md", content="contains hypothesis here"):
        p = os.path.join(tmp, name)
        with open(p, "w") as f:
            f.write(content)
        return {name: ResolvedArtifact(name=name, type="md", path=p, exists=True,
                                       size_bytes=len(content))}

    def test_has_and_contains_pass(self, agent_dir):
        r = self._resolved(agent_dir)
        passed, reasons = eval_rule_gate('has("plan.md") and contains("plan.md", "hypothesis")', r)
        assert passed is True

    def test_contains_fail(self, agent_dir):
        r = self._resolved(agent_dir, content="no keyword")
        passed, _ = eval_rule_gate('contains("plan.md", "hypothesis")', r)
        assert passed is False

    def test_json_len(self, agent_dir):
        p = os.path.join(agent_dir, "refs.json")
        with open(p, "w") as f:
            json.dump(list(range(12)), f)
        r = {"refs.json": ResolvedArtifact("refs.json", "json", p, True, os.path.getsize(p))}
        assert eval_rule_gate('json_len("refs.json") >= 10', r)[0] is True
        assert eval_rule_gate('json_len("refs.json") >= 100', r)[0] is False

    def test_count(self, agent_dir):
        r = self._resolved(agent_dir)
        assert eval_rule_gate("count() >= 1", r)[0] is True

    # --- safety: dangerous expressions must raise GateError, not execute ---

    @pytest.mark.parametrize("expr", [
        '__import__("os").system("echo pwned")',
        'open("/etc/passwd")',
        'has.__class__',
        'count().__class__.__bases__',
        '[].append(1)',
        'lambda: 1',
    ])
    def test_unsafe_expr_rejected(self, expr):
        with pytest.raises(GateError):
            eval_rule_gate(expr, {})


# ============================================================
# 3. Runtime artifact contract validation
# ============================================================

class TestResolveArtifacts:

    def test_missing_required(self, agent_dir):
        st = Stage(id="a", produces=[ArtifactSpec(name="x.md", path="final/x.md", required=True)])
        resolved, problems = resolve_artifacts(st, agent_dir)
        assert len(problems) == 1 and problems[0].startswith("x.md:")
        assert resolved["x.md"].exists is False

    def test_resolved_with_hash(self, agent_dir):
        os.makedirs(os.path.join(agent_dir, "final"))
        with open(os.path.join(agent_dir, "final", "x.md"), "w") as f:
            f.write("data")
        st = Stage(id="a", produces=[ArtifactSpec(name="x.md", path="final/x.md")])
        resolved, missing = resolve_artifacts(st, agent_dir)
        assert missing == []
        assert resolved["x.md"].exists is True
        assert resolved["x.md"].sha256 is not None

    def test_path_breakout_rejected(self, agent_dir):
        """§15 P1-1: absolute / .. paths must not escape the stage workspace."""
        for bad in ("/etc/passwd", "../../../../etc/passwd", "final/../../secret"):
            st = Stage(id="a", produces=[ArtifactSpec(name="x", path=bad, required=True)])
            resolved, problems = resolve_artifacts(st, agent_dir)
            assert problems, f"unsafe path not rejected: {bad}"
            assert resolved["x"].exists is False

    def test_invalid_json_type_fails(self, agent_dir):
        """§15 P2-E: type=json that doesn't parse must NOT count as a valid artifact."""
        os.makedirs(os.path.join(agent_dir, "results"))
        with open(os.path.join(agent_dir, "results", "refs.json"), "w") as f:
            f.write("not valid json {{{")
        st = Stage(id="a", produces=[ArtifactSpec(name="refs.json", type="json",
                                                  path="results/refs.json", required=True)])
        resolved, problems = resolve_artifacts(st, agent_dir)
        assert problems            # invalid json → contract problem
        assert resolved["refs.json"].exists is False


# ============================================================
# 4-5. End-to-end run
# ============================================================

class TestRunPipeline:

    def _two_stage_cfg(self):
        return {"pipeline": {"id": "research", "defaults": {"retry": 1}, "stages": [
            {"id": "plan", "produces": [{"name": "plan.md", "path": "final/plan.md"}],
             "gate": {"kind": "rule", "rule": 'has("plan.md") and contains("plan.md", "hypothesis")'}},
            {"id": "lit", "consumes": ["plan"],
             "produces": [{"name": "refs.json", "path": "results/refs.json", "type": "json"}],
             "gate": {"kind": "rule", "rule": 'json_len("refs.json") >= 3'}},
        ]}}

    def test_completes_and_persists(self, agent_dir):
        execs = {
            "plan": _writer({"final/plan.md": "our hypothesis is X"}),
            "lit": _writer({"results/refs.json": json.dumps([1, 2, 3, 4])}),
        }
        result = run_pipeline(self._two_stage_cfg(), "topic", agent_dir,
                              stage_executor=lambda ec: execs[ec.stage.id](ec))
        assert result.status == PipelineStatus.COMPLETED
        assert [s.status for s in result.stages] == [StageStatus.PASSED, StageStatus.PASSED]
        # verdict + manifest + state persisted
        plan_ws = os.path.join(result.workspace_path, "plan")
        assert os.path.isfile(os.path.join(plan_ws, "verdict.json"))
        assert os.path.isfile(os.path.join(plan_ws, "manifest.json"))
        assert os.path.isfile(os.path.join(result.workspace_path, "pipeline_state.json"))

    def test_multi_upstream_consumes(self, agent_dir):
        """stage 3 consumes two non-adjacent upstream stages via the artifact store."""
        seen = {}

        def s3_exec(ec):
            # read upstream artifacts through the store-backed consumed view
            seen["consumed_keys"] = sorted(ec.consumed.keys())
            seen["plan_path"] = ec.consumed["plan"]["plan.md"].path
            with open(os.path.join(ec.stage_ws, "final", "paper.md"), "w") as f:
                f.write("done")
            return StageExecResult(output="ok", status="completed")

        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "plan", "produces": [{"name": "plan.md", "path": "final/plan.md"}],
             "gate": {"kind": "rule", "rule": 'has("plan.md")'}},
            {"id": "lit", "consumes": ["plan"], "produces": [{"name": "s.md", "path": "final/s.md"}],
             "gate": {"kind": "rule", "rule": 'has("s.md")'}},
            {"id": "paper", "consumes": ["plan", "lit"],
             "produces": [{"name": "paper.md", "path": "final/paper.md"}],
             "gate": {"kind": "rule", "rule": 'has("paper.md")'}},
        ]}}
        execs = {
            "plan": _writer({"final/plan.md": "h"}),
            "lit": _writer({"final/s.md": "survey"}),
            "paper": s3_exec,
        }
        result = run_pipeline(cfg, "x", agent_dir, stage_executor=lambda ec: execs[ec.stage.id](ec))
        assert result.status == PipelineStatus.COMPLETED
        assert seen["consumed_keys"] == ["lit", "plan"]   # both non-adjacent upstreams visible
        assert seen["plan_path"].endswith("plan.md")

    def test_retry_then_halts(self, agent_dir):
        """on_fail=retry + retry=2 → 1+2 attempts, then FAILED/HALTED."""
        calls = {"n": 0}

        def failing(ec):
            calls["n"] += 1
            return StageExecResult(output="", status="completed")  # never writes artifact

        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "retry": 2, "on_fail": "retry",
             "produces": [{"name": "x.md", "path": "final/x.md", "required": True}],
             "gate": {"kind": "rule", "rule": 'has("x.md")'}},
        ]}}
        result = run_pipeline(cfg, "x", agent_dir, stage_executor=failing)
        assert result.status == PipelineStatus.HALTED
        assert result.stages[0].status == StageStatus.FAILED
        assert calls["n"] == 3   # 1 + retry(2)

    def test_halt_does_not_retry(self, agent_dir):
        """§15 P1-4: on_fail=halt ignores retry budget — single attempt, immediate halt."""
        calls = {"n": 0}

        def failing(ec):
            calls["n"] += 1
            return StageExecResult(output="", status="completed")

        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "retry": 2, "on_fail": "halt",   # retry=2 must be ignored
             "produces": [{"name": "x.md", "path": "final/x.md", "required": True}],
             "gate": {"kind": "rule", "rule": 'has("x.md")'}},
        ]}}
        result = run_pipeline(cfg, "x", agent_dir, stage_executor=failing)
        assert result.status == PipelineStatus.HALTED
        assert calls["n"] == 1   # halt = no retry

    def test_skip_on_fail_continues(self, agent_dir):
        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "on_fail": "skip",
             "produces": [{"name": "x.md", "path": "final/x.md", "required": True}],
             "gate": {"kind": "rule", "rule": 'has("x.md")'}},
            {"id": "b", "produces": [{"name": "y.md", "path": "final/y.md"}],
             "gate": {"kind": "rule", "rule": 'has("y.md")'}},
        ]}}
        execs = {"a": _writer({}), "b": _writer({"final/y.md": "ok"})}
        result = run_pipeline(cfg, "x", agent_dir, stage_executor=lambda ec: execs[ec.stage.id](ec))
        assert result.status == PipelineStatus.COMPLETED
        assert result.stages[0].status == StageStatus.SKIPPED
        assert result.stages[1].status == StageStatus.PASSED


# ============================================================
# 6. Resume
# ============================================================

class TestResume:

    def test_resume_skips_passed(self, agent_dir):
        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "produces": [{"name": "x.md", "path": "final/x.md"}],
             "gate": {"kind": "rule", "rule": 'has("x.md")'}},
            {"id": "b", "produces": [{"name": "y.md", "path": "final/y.md"}],
             "gate": {"kind": "rule", "rule": 'has("y.md")'}},
        ]}}

        # First run: stage a passes, stage b halts (writes nothing)
        run_counts = {"a": 0, "b": 0}

        def exec1(ec):
            run_counts[ec.stage.id] += 1
            if ec.stage.id == "a":
                with open(os.path.join(ec.stage_ws, "final", "x.md"), "w") as f:
                    f.write("ok")
            return StageExecResult(output="ok", status="completed")

        cfg["pipeline"]["stages"][1]["on_fail"] = "halt"
        r1 = run_pipeline(cfg, "x", agent_dir, stage_executor=exec1)
        assert r1.status == PipelineStatus.HALTED
        assert run_counts == {"a": 1, "b": 1}

        # Resume: stage a should be skipped (already passed); only b re-runs
        def exec2(ec):
            run_counts[ec.stage.id] += 1
            with open(os.path.join(ec.stage_ws, "final", "y.md"), "w") as f:
                f.write("ok")
            return StageExecResult(output="ok", status="completed")

        r2 = resume_pipeline(r1.workspace_path, cfg, stage_executor=exec2)
        assert r2.status == PipelineStatus.COMPLETED
        assert run_counts["a"] == 1   # NOT re-run
        assert run_counts["b"] == 2   # re-run once on resume

    def test_resume_after_retry_pass_uses_correct_dir(self, agent_dir):
        """§15 P1-2: stage passes on attempt 2 (_a2 dir); downstream after resume must
        consume that attempt's artifact, not the empty base stage dir."""
        # stage a fails attempt 1, passes attempt 2; stage b halts first run, then passes on resume.
        a_calls = {"n": 0}

        def a_exec(ec):
            a_calls["n"] += 1
            if a_calls["n"] >= 2:   # succeed only on the 2nd attempt
                with open(os.path.join(ec.stage_ws, "final", "x.md"), "w") as f:
                    f.write("content-from-attempt-2")
            return StageExecResult(output="ok", status="completed")

        consumed_path = {}

        def b_exec(ec):
            # capture the path a's artifact resolves to through the store
            ra = ec.consumed.get("a", {}).get("x.md")
            consumed_path["p"] = ra.path if ra else None
            return StageExecResult(output="", status="completed")  # fail → halt first time

        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "retry": 2, "on_fail": "retry",
             "produces": [{"name": "x.md", "path": "final/x.md", "required": True}],
             "gate": {"kind": "rule", "rule": 'has("x.md")'}},
            {"id": "b", "consumes": ["a"], "on_fail": "halt",
             "produces": [{"name": "y.md", "path": "final/y.md", "required": True}],
             "gate": {"kind": "rule", "rule": 'has("y.md")'}},
        ]}}

        r1 = run_pipeline(cfg, "x", agent_dir,
                          stage_executor=lambda ec: (a_exec if ec.stage.id == "a" else b_exec)(ec))
        assert r1.status == PipelineStatus.HALTED
        # a passed on attempt 2 → its artifact lives in the _a2 dir
        assert consumed_path["p"] and consumed_path["p"].endswith("x.md")
        assert "_a2" in consumed_path["p"]

        # resume: a skipped (uses persisted _a2 workspace_path); b passes
        def b_exec2(ec):
            ra = ec.consumed.get("a", {}).get("x.md")
            consumed_path["resume"] = ra.path if ra else None
            with open(os.path.join(ec.stage_ws, "final", "y.md"), "w") as f:
                f.write("ok")
            return StageExecResult(output="ok", status="completed")

        r2 = resume_pipeline(r1.workspace_path, cfg,
                             stage_executor=lambda ec: b_exec2(ec))
        assert r2.status == PipelineStatus.COMPLETED
        # after resume, a's artifact still resolves from the correct (_a2) dir
        assert consumed_path["resume"] and "_a2" in consumed_path["resume"]


# ============================================================
# 7. LLM-judge gate
# ============================================================

class TestLLMJudgeGate:

    @pytest.mark.parametrize("text,expected", [
        ('{"score": 0.9, "reasons": ["good"]}', 0.9),
        ('Here is my verdict: {"score":0.4,"reasons":["meh"]} thanks', 0.4),
        ('score: 0.75 overall', 0.75),
        ('1.0', 1.0),
        ('total garbage, no number at all', 0.0),
        ('{"score": 5}', 1.0),     # clamped to [0,1]
    ])
    def test_parse_judge_response(self, text, expected):
        score, _ = _parse_judge_response(text)
        assert abs(score - expected) < 1e-6

    def test_judge_pass_injected(self, agent_dir):
        captured = {}

        def fake_judge(req):
            captured["rubric"] = req.rubric
            captured["text"] = req.artifacts_text
            return JudgeResult(score=0.85, passed=True, reasons=["covers everything"],
                               usage={"provider": "fake", "input_tokens": 10, "output_tokens": 5})

        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "produces": [{"name": "essay.md", "path": "final/essay.md"}],
             "gate": {"kind": "llm_judge", "judge": {"rubric": "states a hypothesis", "threshold": 0.7}}},
        ]}}
        r = run_pipeline(cfg, "x", agent_dir,
                         stage_executor=_writer({"final/essay.md": "our hypothesis is X"}),
                         judge=fake_judge)
        assert r.status == PipelineStatus.COMPLETED
        v = r.stages[0].verdict
        assert v.passed is True and v.score == 0.85
        assert any("score=0.85" in s for s in v.reasons)
        assert any("judge_usage" in s for s in v.reasons)
        assert "hypothesis" in captured["text"]   # judge actually saw the artifact content

    def test_judge_fail_halts(self, agent_dir):
        def low_judge(req):
            return JudgeResult(score=0.30, passed=False, reasons=["too thin"])

        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "on_fail": "halt",
             "produces": [{"name": "essay.md", "path": "final/essay.md"}],
             "gate": {"kind": "llm_judge", "judge": {"rubric": "deep analysis", "threshold": 0.7}}},
        ]}}
        r = run_pipeline(cfg, "x", agent_dir,
                         stage_executor=_writer({"final/essay.md": "thin"}),
                         judge=low_judge)
        assert r.status == PipelineStatus.HALTED
        assert r.stages[0].status == StageStatus.FAILED
        assert r.stages[0].verdict.score == 0.30

    @pytest.mark.skipif(not _ollama_up(), reason="ollama not reachable on :11434")
    def test_judge_live_ollama(self, agent_dir):
        """Real local-ollama judge — backs the self-hosted '逐级验收' selling point."""
        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a", "on_fail": "halt",
             "produces": [{"name": "essay.md", "path": "final/essay.md"}],
             "gate": {"kind": "llm_judge", "judge": {
                 "model": {"provider": "ollama", "name": "qwen2.5:7b"},
                 "rubric": "The text must clearly state a research hypothesis.",
                 "threshold": 0.5}}},
        ]}}
        good = ("Our central hypothesis is that cost-graded effect types can predict "
                "an agent's resource consumption before execution. We test it on 835 configs.")
        r = run_pipeline(cfg, "topic", agent_dir,
                         stage_executor=_writer({"final/essay.md": good}))
        v = r.stages[0].verdict
        assert v is not None and v.score is not None
        assert 0.0 <= v.score <= 1.0
        assert r.status in (PipelineStatus.COMPLETED, PipelineStatus.HALTED)


# ============================================================
# 8. Cost estimation (②: judge as term, Paper34 prediction side)
# ============================================================

class TestCostEstimate:

    def test_estimate_pipeline_cost_includes_judge(self):
        """estimate_pipeline_cost folds agent + judge-term cost; llm_judge stage > 0."""
        cfg = {"pipeline": {"id": "p", "stages": [
            {"id": "a",
             "agent": {"type": "simple",
                       "model": {"provider": "ollama", "name": "qwen2.5:7b"},
                       "systemPrompt": "x"},
             "produces": [{"name": "a.md", "path": "final/a.md"}],
             "gate": {"kind": "llm_judge",
                      "judge": {"model": {"provider": "ollama", "name": "qwen2.5:7b"},
                                "rubric": "r", "threshold": 0.5}}},
        ]}}
        total, breakdown = estimate_pipeline_cost(load_pipeline(cfg))
        assert len(breakdown) == 1
        assert breakdown[0][0] == "a"
        assert total.tokens > 0            # judge term contributes a token upper bound
        assert 0 < total.probability <= 1

    def test_build_judge_term_is_a_term(self):
        from lambdagent.cost_grade import estimate_cost
        term = build_judge_term({"rubric": "r", "model": {"provider": "ollama", "name": "qwen2.5:7b"}})
        g = estimate_cost(term)            # judge is in the term graph → cost computable
        assert g.tokens > 0


# ============================================================
# 9. Live ollama: term judge + research demo end-to-end
# ============================================================

class TestLiveOllama:

    @pytest.mark.skipif(not _ollama_up(), reason="ollama not reachable on :11434")
    def test_term_judge_live(self, agent_dir):
        req = JudgeRequest(
            rubric="The text must clearly state a research hypothesis.",
            threshold=0.5, model={"provider": "ollama", "name": "qwen2.5:7b"},
            artifacts_text="### x\nOur hypothesis is that graded cost predicts agent resource use.")
        jr = term_judge(req)
        assert 0.0 <= jr.score <= 1.0
        assert jr.usage and jr.usage.get("in_term_graph") is True

    @pytest.mark.skipif(not _ollama_up(), reason="ollama not reachable on :11434")
    def test_research_demo_end_to_end(self, agent_dir):
        """① Real plan→lit run on local ollama via the research-demo preset.

        Asserts ORCHESTRATION invariants (terminal status, stage order, state +
        verdict persisted), not the small model's exact output — a live local-LLM
        e2e must not flake on model nondeterminism or transient ollama hiccups.
        """
        assert os.path.isfile(_DEMO_PRESET)
        r = run_pipeline(_DEMO_PRESET, "形式化资源语义在 AI agent 调度中的应用", agent_dir)
        # orchestration reached a terminal state and ran in declared order
        assert r.status in (PipelineStatus.COMPLETED, PipelineStatus.HALTED)
        assert [s.stage_id for s in r.stages][:1] == ["plan"]
        assert os.path.isfile(os.path.join(r.workspace_path, "pipeline_state.json"))
        # when the plan stage passes, the artifact bridge must have produced plan.md
        plan = r.stages[0]
        if plan.status == StageStatus.PASSED:
            assert os.path.isfile(os.path.join(plan.workspace_path, "final", "plan.md"))


# ============================================================
# Human gate: YIELD + resolve_human_gate（§6.3，P3 完成版）
# ============================================================

from agentpaas.engine.pipeline import resolve_human_gate


class TestHumanGate:

    def _cfg(self, on_fail="halt"):
        """s1 规则门 → s2 人工门 → s3 规则门（验证断点前后都正确）。"""
        return {"pipeline": {"id": "hg", "stages": [
            {"id": "s1",
             "produces": [{"name": "plan", "type": "md", "path": "plan.md"}],
             "gate": {"kind": "rule", "rule": 'has("plan")'}},
            {"id": "s2", "consumes": ["s1"],
             "produces": [{"name": "draft", "type": "md", "path": "draft.md"}],
             "gate": {"kind": "human", "prompt": "请导师确认草稿方向"},
             "on_fail": on_fail, "retry": 1},
            {"id": "s3", "consumes": ["s2"],
             "produces": [{"name": "final", "type": "md", "path": "final.md"}],
             "gate": {"kind": "rule", "rule": 'has("final")'}},
        ]}}

    def _execs(self, counter):
        def make(files):
            def _exec(ec):
                counter[ec.stage.id] = counter.get(ec.stage.id, 0) + 1
                for rel, content in files.items():
                    p = os.path.join(ec.stage_ws, rel)
                    os.makedirs(os.path.dirname(p) or ec.stage_ws, exist_ok=True)
                    with open(p, "w", encoding="utf-8") as f:
                        f.write(content)
                return StageExecResult(output="ok", status="completed")
            return _exec
        table = {"s1": make({"plan.md": "the plan"}),
                 "s2": make({"draft.md": "the draft"}),
                 "s3": make({"final.md": "the final"})}
        return lambda ec: table[ec.stage.id](ec)

    def test_yields_at_human_gate(self, agent_dir):
        n = {}
        r = run_pipeline(self._cfg(), "topic", agent_dir, stage_executor=self._execs(n))
        assert r.status == PipelineStatus.AWAITING_HUMAN
        assert n == {"s1": 1, "s2": 1}  # s3 没跑
        # state 落盘正确，verdict 是 pending 占位
        state = json.load(open(os.path.join(r.workspace_path, "pipeline_state.json")))
        s2 = state["stages"]["s2"]
        assert s2["status"] == "awaiting_human"
        assert s2["gate_prompt"] == "请导师确认草稿方向"
        v = json.load(open(os.path.join(s2["workspace_path"], "verdict.json")))
        assert v["passed"] is False and "PENDING_HUMAN" in v["reasons"][0]

    def test_approved_resumes_to_completion(self, agent_dir):
        n = {}
        ex = self._execs(n)
        r = run_pipeline(self._cfg(), "topic", agent_dir, stage_executor=ex)
        r2 = resolve_human_gate(self._cfg(), r.workspace_path, "approved",
                                note="方向没问题", stage_executor=ex)
        assert r2.status == PipelineStatus.COMPLETED
        # s1/s2 没重算，只补跑了 s3
        assert n == {"s1": 1, "s2": 1, "s3": 1}
        # human verdict 覆盖了 pending 占位
        state = json.load(open(os.path.join(r.workspace_path, "pipeline_state.json")))
        s2 = state["stages"]["s2"]
        assert s2["status"] == "passed"
        v = json.load(open(s2["verdict_path"]))
        assert v["passed"] is True and "方向没问题" in v["reasons"][0]

    def test_rejected_halt(self, agent_dir):
        n = {}
        ex = self._execs(n)
        r = run_pipeline(self._cfg(on_fail="halt"), "topic", agent_dir, stage_executor=ex)
        r2 = resolve_human_gate(self._cfg(on_fail="halt"), r.workspace_path,
                                "rejected", note="重写", stage_executor=ex)
        assert r2.status == PipelineStatus.HALTED
        assert n == {"s1": 1, "s2": 1}  # 不重跑
        state = json.load(open(os.path.join(r.workspace_path, "pipeline_state.json")))
        assert state["stages"]["s2"]["status"] == "failed"

    def test_rejected_retry_redoes_stage_then_awaits_again(self, agent_dir):
        n = {}
        ex = self._execs(n)
        r = run_pipeline(self._cfg(on_fail="retry"), "topic", agent_dir, stage_executor=ex)
        r2 = resolve_human_gate(self._cfg(on_fail="retry"), r.workspace_path,
                                "rejected", note="再来一版", stage_executor=ex)
        # 重做后再次停在人工门
        assert r2.status == PipelineStatus.AWAITING_HUMAN
        assert n == {"s1": 1, "s2": 2}  # s2 重跑一次, s1 不重算
        # 再批准 → 跑完
        r3 = resolve_human_gate(self._cfg(on_fail="retry"), r.workspace_path,
                                "approved", stage_executor=ex)
        assert r3.status == PipelineStatus.COMPLETED
        assert n == {"s1": 1, "s2": 2, "s3": 1}

    def test_rejected_skip_continues(self, agent_dir):
        n = {}
        ex = self._execs(n)
        r = run_pipeline(self._cfg(on_fail="skip"), "topic", agent_dir, stage_executor=ex)
        r2 = resolve_human_gate(self._cfg(on_fail="skip"), r.workspace_path,
                                "rejected", stage_executor=ex)
        assert r2.status == PipelineStatus.COMPLETED
        state = json.load(open(os.path.join(r.workspace_path, "pipeline_state.json")))
        assert state["stages"]["s2"]["status"] == "skipped"
        assert n == {"s1": 1, "s2": 1, "s3": 1}

    def test_bad_decision_and_no_awaiting(self, agent_dir):
        n = {}
        ex = self._execs(n)
        r = run_pipeline(self._cfg(), "topic", agent_dir, stage_executor=ex)
        with pytest.raises(PipelineError):
            resolve_human_gate(self._cfg(), r.workspace_path, "maybe")
        resolve_human_gate(self._cfg(), r.workspace_path, "approved", stage_executor=ex)
        with pytest.raises(PipelineError):  # 已无 awaiting 阶段
            resolve_human_gate(self._cfg(), r.workspace_path, "approved", stage_executor=ex)