gyxx-flow/tests/test_langgraph_workflow_engine.py

from __future__ import annotations

import json
import threading
from collections import deque
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from langgraph.graph.state import CompiledStateGraph

from gyxx_flow.catalog import WorkflowCatalog, WorkflowEntry
from gyxx_flow.cli import build_default_registry
from gyxx_flow.core.config import Settings
from gyxx_flow.core.context import RunContext
from gyxx_flow.core.layout import DataLayout
from gyxx_flow.core.locks import LockManager
from gyxx_flow.core.records import RunJournal
from gyxx_flow.modules.shop_intelligence import ShopIntelligenceModule
from gyxx_flow.ops.effects import EffectLedger
from gyxx_flow.workflow.engine import WorkflowEngine
from gyxx_flow.workflow.graph import (
    GraphStepResult,
    WorkflowGraphState,
    compile_workflow_graph,
    initial_workflow_graph_state,
)
from gyxx_flow.workflow.model import StepDefinition, WorkflowDefinition
from gyxx_flow.workflow.steps import StepExecution


class FakeStep:
    def __init__(self, *outcomes: StepExecution) -> None:
        self.outcomes = deque(outcomes or (StepExecution(exit_code=0),))
        self.calls: list[dict[str, Any]] = []

    def execute(
        self,
        *,
        context: RunContext,
        timeout_seconds: float | None,
        dry_run: bool,
    ) -> StepExecution:
        self.calls.append(
            {
                "context": context,
                "timeout_seconds": timeout_seconds,
                "dry_run": dry_run,
            }
        )
        return self.outcomes.popleft()


class CoordinatedStep:
    def __init__(
        self,
        barrier: threading.Barrier,
        *,
        error: str | None = None,
    ) -> None:
        self.barrier = barrier
        self.error = error
        self.calls = 0

    def execute(
        self,
        *,
        context: RunContext,
        timeout_seconds: float | None,
        dry_run: bool,
    ) -> StepExecution:
        del context, timeout_seconds, dry_run
        self.calls += 1
        self.barrier.wait(timeout=3)
        if self.error is not None:
            raise RuntimeError(self.error)
        return StepExecution(exit_code=0)


def _context(*, suffix: str = "abc123", shadow: bool = False) -> RunContext:
    return RunContext.create(
        "shop.weekly",
        "2026-07-27",
        shadow=shadow,
        now=datetime(2026, 7, 27, 4, 0, tzinfo=timezone.utc),
        random_suffix=suffix,
    )


def _success_state(
    step: StepDefinition,
    state: WorkflowGraphState,
) -> GraphStepResult:
    del state
    return {
        "step_id": step.step_id,
        "status": "success",
        "attempts": [
            {
                "attempt": 1,
                "status": "success",
                "exit_code": 0,
                "error": None,
                "reason": None,
            }
        ],
        "final_exit_code": 0,
        "critical": step.critical,
        "reason": None,
    }


def test_compiler_creates_one_langgraph_node_per_step_in_stable_order() -> None:
    workflow = WorkflowDefinition(
        "shop.weekly",
        (
            StepDefinition("publish", FakeStep(), depends_on=("collect",)),
            StepDefinition("audit", FakeStep()),
            StepDefinition("collect", FakeStep()),
        ),
    )

    compiled = compile_workflow_graph(workflow, run_step=_success_state)
    graph = compiled.get_graph()

    assert isinstance(compiled, CompiledStateGraph)
    assert set(graph.nodes) == {"__start__", "audit", "collect", "publish", "__end__"}
    assert {(edge.source, edge.target) for edge in graph.edges} == {
        ("__start__", "audit"),
        ("__start__", "collect"),
        ("audit", "__end__"),
        ("collect", "publish"),
        ("publish", "__end__"),
    }

    state = initial_workflow_graph_state(_context(), dry_run=False)
    final_state = compiled.invoke(state)

    assert set(final_state["step_results"]) == {"audit", "collect", "publish"}
    json.dumps(final_state)


def test_compiler_joins_uneven_dependency_branches_before_running_once() -> None:
    workflow = WorkflowDefinition(
        "shop.weekly",
        (
            StepDefinition("source", FakeStep()),
            StepDefinition("parallel", FakeStep()),
            StepDefinition("derived", FakeStep(), depends_on=("source",)),
            StepDefinition(
                "publish",
                FakeStep(),
                depends_on=("parallel", "derived"),
            ),
        ),
    )
    publish_inputs: list[set[str]] = []
    calls: list[str] = []

    def run_step(
        step: StepDefinition,
        state: WorkflowGraphState,
    ) -> GraphStepResult:
        calls.append(step.step_id)
        if step.step_id == "publish":
            publish_inputs.append(set(state["step_results"]))
        return _success_state(step, state)

    compiled = compile_workflow_graph(workflow, run_step=run_step)
    compiled.invoke(initial_workflow_graph_state(_context(), dry_run=False))

    assert calls.count("publish") == 1
    assert publish_inputs == [{"source", "parallel", "derived"}]


def test_engine_preserves_retry_dependency_and_journal_contract(tmp_path: Path) -> None:
    flaky = FakeStep(
        StepExecution(exit_code=7, error="temporary"),
        StepExecution(exit_code=0),
    )
    downstream = FakeStep()
    workflow = WorkflowDefinition(
        "shop.weekly",
        (
            StepDefinition("flaky", flaky, max_attempts=2),
            StepDefinition("required", FakeStep(StepExecution(exit_code=9, error="bad"))),
            StepDefinition("downstream", downstream, depends_on=("required",)),
        ),
    )
    context = _context()
    journal = RunJournal.create(DataLayout(tmp_path), context)

    result = WorkflowEngine(LockManager(tmp_path / "locks")).execute(
        workflow,
        context=context,
        journal=journal,
    )

    assert result.status == "failed"
    assert [attempt.exit_code for attempt in result.steps["flaky"].attempts] == [7, 0]
    assert result.steps["downstream"].reason == "dependency-failed"
    assert downstream.calls == []
    payload = json.loads(journal.path.read_text(encoding="utf-8"))
    assert payload["steps"]["flaky.attempt-1"]["status"] == "failed"
    assert payload["steps"]["flaky.attempt-2"]["status"] == "success"
    assert payload["steps"]["downstream.attempt-1"]["status"] == "skipped"


def test_independent_graph_nodes_run_concurrently_and_isolate_action_failure(
    tmp_path: Path,
) -> None:
    barrier = threading.Barrier(2)
    jd = CoordinatedStep(barrier, error="JD browser session failed")
    tmall = CoordinatedStep(barrier)
    workflow = WorkflowDefinition(
        "shop.weekly",
        (
            StepDefinition("jd", jd, resources=("browser:jd",)),
            StepDefinition("tmall", tmall, resources=("browser:tmall",)),
        ),
    )
    context = _context(suffix="parallel1")
    journal = RunJournal.create(DataLayout(tmp_path), context)

    result = WorkflowEngine(LockManager(tmp_path / "locks")).execute(
        workflow,
        context=context,
        journal=journal,
    )

    assert jd.calls == 1
    assert tmall.calls == 1
    assert result.status == "failed"
    assert result.steps["jd"].status == "failed"
    assert result.steps["tmall"].status == "success"
    payload = json.loads(journal.path.read_text(encoding="utf-8"))
    assert payload["error"] == (
        "critical steps failed: jd: RuntimeError: JD browser session failed"
    )
    assert payload["steps"]["jd.attempt-1"]["error"] == (
        "RuntimeError: JD browser session failed"
    )
    assert payload["steps"]["tmall.attempt-1"]["status"] == "success"


def test_workflow_error_reports_each_failed_step_with_bounded_details(
    tmp_path: Path,
) -> None:
    jd_error = "JD PG write refused: " + "x" * 5_000 + " JD_ROOT_CAUSE_AT_TAIL"
    tmall_error = (
        "TM Feishu insert refused: "
        + "y" * 5_000
        + " TM_ROOT_CAUSE_AT_TAIL"
    )
    workflow = WorkflowDefinition(
        "shop.weekly",
        (
            StepDefinition("jd", FakeStep(StepExecution(exit_code=7, error=jd_error))),
            StepDefinition(
                "tmall",
                FakeStep(StepExecution(exit_code=9, error=tmall_error)),
            ),
        ),
    )
    context = _context(suffix="bounded1")
    journal = RunJournal.create(DataLayout(tmp_path), context)

    WorkflowEngine(LockManager(tmp_path / "locks")).execute(
        workflow,
        context=context,
        journal=journal,
    )

    error = json.loads(journal.path.read_text(encoding="utf-8"))["error"]
    assert "jd: JD PG write refused" in error
    assert "tmall: TM Feishu insert refused" in error
    assert "JD_ROOT_CAUSE_AT_TAIL" in error
    assert "TM_ROOT_CAUSE_AT_TAIL" in error
    assert len(error) <= 4_096


def test_engine_preserves_dry_run_and_shadow_policy(tmp_path: Path) -> None:
    dry_action = FakeStep()
    dry_workflow = WorkflowDefinition(
        "shop.weekly", (StepDefinition("collect", dry_action),)
    )
    dry_context = _context()
    dry_result = WorkflowEngine(LockManager(tmp_path / "locks")).execute(
        dry_workflow,
        context=dry_context,
        journal=RunJournal.create(DataLayout(tmp_path), dry_context),
        dry_run=True,
    )

    sink = FakeStep()
    shadow_workflow = WorkflowDefinition(
        "shop.weekly",
        (StepDefinition("publish", sink, production_sink=True),),
    )
    shadow_context = _context(suffix="shadow1", shadow=True)
    shadow_result = WorkflowEngine(LockManager(tmp_path / "locks")).execute(
        shadow_workflow,
        context=shadow_context,
        journal=RunJournal.create(DataLayout(tmp_path), shadow_context),
    )

    assert dry_result.steps["collect"].reason == "dry-run"
    assert shadow_result.steps["publish"].reason == "shadow-policy"
    assert dry_action.calls == []
    assert sink.calls == []


def test_engine_preserves_effect_ledger_replay_policy(tmp_path: Path) -> None:
    action = FakeStep()
    workflow = WorkflowDefinition(
        "shop.weekly",
        (StepDefinition("publish", action, production_sink=True),),
    )
    engine = WorkflowEngine(
        LockManager(tmp_path / "locks"),
        effect_ledger=EffectLedger(tmp_path),
    )
    first_context = _context(suffix="first1")
    second_context = _context(suffix="second2")

    first = engine.execute(
        workflow,
        context=first_context,
        journal=RunJournal.create(DataLayout(tmp_path), first_context),
    )
    second = engine.execute(
        workflow,
        context=second_context,
        journal=RunJournal.create(DataLayout(tmp_path), second_context),
    )

    assert first.steps["publish"].status == "success"
    assert second.steps["publish"].reason == "idempotency-replay"
    assert len(action.calls) == 1


def test_every_registered_workflow_compiles_to_langgraph(tmp_path: Path) -> None:
    project_root = Path(__file__).parents[1]
    registry = build_default_registry(
        Settings(project_root=project_root, data_root=tmp_path)
    )

    compiled_count = 0
    for entry in registry.catalog.workflows:
        if not registry.is_registered(entry.workflow_id):
            continue
        definition = registry.resolve(entry.workflow_id).definition
        compiled = compile_workflow_graph(definition, run_step=_success_state)
        assert isinstance(compiled, CompiledStateGraph)
        compiled_count += 1

    assert compiled_count > 0


def test_jd_self_operated_uses_source_timeouts_and_continues_after_brand_failure(
    tmp_path: Path,
) -> None:
    calls: list[tuple[str, float | None]] = []

    class RecordedCommand:
        def __init__(self, entry: WorkflowEntry) -> None:
            self.entry = entry

        def execute(
            self,
            *,
            context: RunContext,
            timeout_seconds: float | None,
            dry_run: bool,
        ) -> StepExecution:
            del context
            assert dry_run is False
            calls.append((self.entry.entry, timeout_seconds))
            if self.entry.entry.endswith("jd_self_operated_brand_daily.py"):
                return StepExecution(exit_code=3, error="brand failed")
            return StepExecution(exit_code=0)

    def command_factory(
        entry: WorkflowEntry,
        _context: RunContext,
    ) -> RecordedCommand:
        return RecordedCommand(entry)

    catalog = WorkflowCatalog.load(Path(__file__).parents[1] / "config")
    workflow = next(
        item
        for item in ShopIntelligenceModule.from_catalog(
            catalog,
            command_factory=command_factory,
        ).workflow_definitions()
        if item.workflow_id == "shop.jd_self_operated.daily"
    )
    context = RunContext.create(
        workflow.workflow_id,
        "2026-08-02",
        random_suffix="jdself1",
    )

    result = WorkflowEngine(LockManager(tmp_path / "locks")).execute(
        workflow,
        context=context,
        journal=RunJournal.create(DataLayout(tmp_path), context),
    )

    assert calls == [
        ("collectors/jd_self_operated_brand_daily.py", 600),
        ("collectors/jd_self_operated_product_daily.py", 1_800),
    ]
    assert list(result.steps) == ["brand", "product"]
    assert result.steps["brand"].status == "failed"
    assert result.steps["product"].status == "success"
    assert result.status == "failed"
    assert result.exit_code == 1