宜博 / lemongraph

"""U-011 tests: c_work_graph_plan, mocked at the subprocess boundary (CLAUDE.md
Testing constraint). Real-call coverage is in src/test_work_graph_plan_e2e.py
(RUN_REAL_AGENT_E2E=1) -- that real call is what proved --json-schema doesn't
constrain output when --agent is set, which is why these mocks return
free-form ```json-fenced text in `result`, not `structured_output`.

Run: python3 -m unittest src.test_work_graph_plan -v
"""

import json
import subprocess
import unittest
from unittest.mock import patch

from src import schema, work_graph_execute, work_graph_orchestrate, work_graph_plan

_MANIFEST = {
    "nodes": [
        {"node_id": "n1", "artifact_type": "brd", "depends_on": []},
        {"node_id": "n2", "artifact_type": "prd", "depends_on": ["n1"]},
        {"node_id": "n3", "artifact_type": "sdd", "depends_on": ["n2"]},
        {"node_id": "n4", "artifact_type": "benchmark", "depends_on": ["n3"]},
    ],
    "edges": [
        {"from": "n1", "to": "n2", "edge_type": "needs_judgment"},
        {"from": "n2", "to": "n3", "edge_type": "needs_judgment"},
        {"from": "n3", "to": "n4", "edge_type": "needs_judgment"},
    ],
}


def _completed_process(returncode=0, stdout_obj=None, stderr=""):
    stdout = json.dumps(stdout_obj) if stdout_obj is not None else ""
    return subprocess.CompletedProcess(args=["claude"], returncode=returncode, stdout=stdout, stderr=stderr)


def _fenced_result(manifest, wrapper_key=None):
    body = {wrapper_key: manifest} if wrapper_key else manifest
    return "Here is my plan:\n\n```json\n" + json.dumps(body) + "\n```\n\nDone."


class PlanWorkGraphTestBase(unittest.TestCase):
    def setUp(self):
        # check_same_thread=False: GenericPathEndToEndViaRunGraphTest below
        # drives a real multi-node concurrent layer through run_graph()
        # (G), which calls into this conn from ThreadPoolExecutor worker
        # threads via _LockedConn -- see test_work_graph_orchestrate.py's
        # RunGraphTestBase for the same reasoning.
        self.conn = schema.connect(":memory:", check_same_thread=False)
        schema.create_schema(self.conn)
        self.conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
        )
        self.conn.execute(
            "INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) "
            "VALUES ('task-1', 'proj-1', '实现贪吃蛇游戏', 'goal', 'queued', 'planner')"
        )
        self.conn.commit()

    def tearDown(self):
        self.conn.close()


class ExtractGraphManifestTest(unittest.TestCase):
    def test_bare_json_result(self):
        self.assertEqual(work_graph_plan._extract_graph_manifest(json.dumps(_MANIFEST)), _MANIFEST)

    def test_fenced_json_block_with_prose_around_it(self):
        self.assertEqual(work_graph_plan._extract_graph_manifest(_fenced_result(_MANIFEST)), _MANIFEST)

    def test_fenced_json_wrapped_under_graph_manifest_key(self):
        text = _fenced_result(_MANIFEST, wrapper_key="graph_manifest")
        self.assertEqual(work_graph_plan._extract_graph_manifest(text), _MANIFEST)

    def test_no_json_anywhere_returns_none(self):
        self.assertIsNone(work_graph_plan._extract_graph_manifest("sorry, I can't do that"))


class NormalizeGraphManifestTest(unittest.TestCase):
    """U-031 regression: a real Planner call (2026-08-01) came back using
    `id` instead of the documented `node_id`, and never populated any
    node's own `depends_on` -- dependencies existed only as `edges`. This
    is a real, observed deviation, not a hypothetical one; the fixture
    below is the actual shape captured from that call (fields trimmed to
    what normalization touches)."""

    def test_already_well_formed_manifest_untouched(self):
        manifest = {
            "nodes": [
                {"node_id": "brd", "artifact_type": "brd", "depends_on": []},
                {"node_id": "prd", "artifact_type": "prd", "depends_on": ["brd"]},
            ],
            "edges": [{"from": "brd", "to": "prd", "edge_type": "needs_judgment"}],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual(result["nodes"][1]["depends_on"], ["brd"])

    def test_id_key_normalized_to_node_id(self):
        manifest = {
            "nodes": [{"id": "brd", "artifact_type": "brd"}, {"id": "prd", "artifact_type": "prd"}],
            "edges": [{"from": "brd", "to": "prd", "edge_type": "needs_judgment"}],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual([n["node_id"] for n in result["nodes"]], ["brd", "prd"])

    def test_depends_on_derived_from_edges_when_absent(self):
        manifest = {
            "nodes": [{"id": "brd"}, {"id": "prd"}, {"id": "sdd"}, {"id": "benchmark"}],
            "edges": [
                {"from": "brd", "to": "prd", "edge_type": "needs_judgment"},
                {"from": "prd", "to": "sdd", "edge_type": "needs_judgment"},
                {"from": "sdd", "to": "benchmark", "edge_type": "needs_judgment"},
            ],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        by_id = {n["node_id"]: n for n in result["nodes"]}
        self.assertEqual(by_id["brd"]["depends_on"], [])
        self.assertEqual(by_id["prd"]["depends_on"], ["brd"])
        self.assertEqual(by_id["sdd"]["depends_on"], ["prd"])
        self.assertEqual(by_id["benchmark"]["depends_on"], ["sdd"])

    def test_does_not_duplicate_dependency_already_present(self):
        manifest = {
            "nodes": [{"node_id": "brd", "depends_on": []}, {"node_id": "prd", "depends_on": ["brd"]}],
            "edges": [{"from": "brd", "to": "prd", "edge_type": "needs_judgment"}],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual(result["nodes"][1]["depends_on"], ["brd"])

    def test_edge_referencing_unknown_node_id_ignored(self):
        manifest = {"nodes": [{"id": "brd"}], "edges": [{"from": "brd", "to": "no-such-node"}]}
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual(result["nodes"][0]["depends_on"], [])

    def test_zero_dependency_info_leaves_nodes_independent(self):
        # 2026-08-08 (FR-B4 "全部自由规划"): the canonical-doc-chain forced-
        # order enforcement that used to backfill this exact case (2026-08-03
        # incident) is deliberately removed, not extended -- Planner's node
        # set is no longer a fixed architectural constant. A real response
        # with zero dependency info now genuinely produces zero depends_on,
        # same as any other node set; robustness rests entirely on the
        # edge-derived backfill (test_depends_on_derived_from_edges_when_absent)
        # plus _validate_loop_back_targets, not a hardcoded chain to guess from.
        manifest = {
            "nodes": [
                {"node_id": "brd", "artifact_type": "brd"},
                {"node_id": "prd", "artifact_type": "prd"},
                {"node_id": "sdd", "artifact_type": "sdd"},
                {"node_id": "benchmark", "artifact_type": "benchmark"},
                {"node_id": "test_dataset", "artifact_type": "test_dataset"},
            ],
            "edges": [],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        for node in result["nodes"]:
            self.assertEqual(node["depends_on"], [])

    def test_on_failure_edge_does_not_get_folded_into_depends_on(self):
        # 2026-08-08: an on_failure edge is control-flow, not a data
        # dependency -- the edge-derived depends_on backfill above must
        # skip it (same reasoning as work_graph_orchestrate.py's
        # topological_layers()/_transitive_dependents).
        manifest = {
            "nodes": [
                {"node_id": "sdd", "artifact_type": "sdd"},
                {
                    "node_id": "benchmark", "artifact_type": "benchmark", "depends_on": ["sdd"],
                    "on_fail": {"action": "loop_back"},
                },
            ],
            "edges": [{"from": "benchmark", "to": "sdd", "edge_type": "on_failure"}],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        by_id = {n["node_id"]: n for n in result["nodes"]}
        self.assertEqual(by_id["sdd"]["depends_on"], [])


class TransitiveAncestorsTest(unittest.TestCase):
    def test_includes_self_and_all_upstream(self):
        manifest = {
            "nodes": [
                {"node_id": "a", "depends_on": []},
                {"node_id": "b", "depends_on": ["a"]},
                {"node_id": "c", "depends_on": ["b"]},
            ],
            "edges": [],
        }
        self.assertEqual(work_graph_plan._transitive_ancestors(manifest, "c"), {"a", "b", "c"})
        self.assertEqual(work_graph_plan._transitive_ancestors(manifest, "a"), {"a"})

    def test_ignores_on_failure_edges(self):
        manifest = {
            "nodes": [{"node_id": "sdd"}, {"node_id": "benchmark"}],
            "edges": [{"from": "benchmark", "to": "sdd", "edge_type": "on_failure"}],
        }
        # an on_failure edge alone (no real depends_on/data edge) must NOT
        # make benchmark an ancestor of sdd.
        self.assertEqual(work_graph_plan._transitive_ancestors(manifest, "sdd"), {"sdd"})


class ValidateLoopBackTargetsTest(unittest.TestCase):
    def test_valid_loop_back_passes(self):
        manifest = {
            "nodes": [
                {"node_id": "sdd", "depends_on": []},
                {
                    "node_id": "benchmark", "depends_on": ["sdd"],
                    "on_fail": {"action": "loop_back"},
                },
            ],
            "edges": [{"from": "benchmark", "to": "sdd", "edge_type": "on_failure"}],
        }
        work_graph_plan._validate_loop_back_targets(manifest)  # must not raise

    def test_no_on_failure_edge_raises(self):
        manifest = {
            "nodes": [
                {"node_id": "sdd", "depends_on": []},
                {"node_id": "benchmark", "depends_on": ["sdd"], "on_fail": {"action": "loop_back"}},
            ],
            "edges": [],
        }
        with self.assertRaises(work_graph_plan.PlanningFailedError):
            work_graph_plan._validate_loop_back_targets(manifest)

    def test_two_on_failure_edges_from_same_node_raises(self):
        manifest = {
            "nodes": [
                {"node_id": "sdd", "depends_on": []},
                {"node_id": "prd", "depends_on": []},
                {
                    "node_id": "benchmark", "depends_on": ["sdd", "prd"],
                    "on_fail": {"action": "loop_back"},
                },
            ],
            "edges": [
                {"from": "benchmark", "to": "sdd", "edge_type": "on_failure"},
                {"from": "benchmark", "to": "prd", "edge_type": "on_failure"},
            ],
        }
        with self.assertRaises(work_graph_plan.PlanningFailedError):
            work_graph_plan._validate_loop_back_targets(manifest)

    def test_target_not_an_upstream_dependency_raises(self):
        # sdd and unrelated are siblings with no dependency relationship --
        # benchmark depends on sdd but its on_failure edge points at the
        # unrelated node instead, which never influenced benchmark at all.
        manifest = {
            "nodes": [
                {"node_id": "sdd", "depends_on": []},
                {"node_id": "unrelated", "depends_on": []},
                {
                    "node_id": "benchmark", "depends_on": ["sdd"],
                    "on_fail": {"action": "loop_back"},
                },
            ],
            "edges": [{"from": "benchmark", "to": "unrelated", "edge_type": "on_failure"}],
        }
        with self.assertRaises(work_graph_plan.PlanningFailedError):
            work_graph_plan._validate_loop_back_targets(manifest)

    def test_target_not_in_manifest_at_all_raises(self):
        manifest = {
            "nodes": [{"node_id": "benchmark", "depends_on": [], "on_fail": {"action": "loop_back"}}],
            "edges": [{"from": "benchmark", "to": "nonexistent", "edge_type": "on_failure"}],
        }
        with self.assertRaises(work_graph_plan.PlanningFailedError):
            work_graph_plan._validate_loop_back_targets(manifest)

    def test_nodes_without_loop_back_are_unaffected(self):
        manifest = {
            "nodes": [
                {"node_id": "brd", "depends_on": []},
                {"node_id": "prd", "depends_on": ["brd"]},
            ],
            "edges": [],
        }
        work_graph_plan._validate_loop_back_targets(manifest)  # must not raise

    def test_normalize_graph_manifest_calls_validation_and_raises_on_bad_loop_back(self):
        manifest = {
            "nodes": [{"node_id": "benchmark", "on_fail": {"action": "loop_back"}}],
            "edges": [],
        }
        with self.assertRaises(work_graph_plan.PlanningFailedError):
            work_graph_plan._normalize_graph_manifest(manifest)


class HappyPathTest(PlanWorkGraphTestBase):
    def test_returns_run_id_and_graph_manifest(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            result = work_graph_plan.plan_work_graph(
                self.conn, "planner", "task-1", "实现贪吃蛇游戏", "/tmp/proj-1"
            )
        self.assertEqual(result["graph_manifest"], _MANIFEST)
        self.assertTrue(result["run_id"])

    def test_creates_work_graph_run_row_in_planning_status(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            result = work_graph_plan.plan_work_graph(
                self.conn, "planner", "task-1", "实现贪吃蛇游戏", "/tmp/proj-1"
            )
        row = self.conn.execute(
            "SELECT task_id, org_node, status, graph_manifest FROM work_graph_run WHERE run_id = ?",
            (result["run_id"],),
        ).fetchone()
        self.assertEqual(row[0], "task-1")
        self.assertEqual(row[1], "planner")
        self.assertEqual(row[2], "planning")
        self.assertEqual(json.loads(row[3]), _MANIFEST)

    def test_invokes_claude_with_expected_args(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")

        args, kwargs = run.call_args
        argv = args[0]
        self.assertEqual(argv[:3], ["claude", "--agent", "planner"])
        self.assertIn("--permission-mode", argv)
        self.assertIn("auto", argv)
        self.assertNotIn("--bare", argv)
        # --json-schema deliberately dropped -- confirmed for real it's a no-op
        # when combined with --agent (see module docstring)
        self.assertNotIn("--json-schema", argv)
        self.assertEqual(kwargs["cwd"], "/tmp/proj-1")
        self.assertEqual(kwargs["stdin"], subprocess.DEVNULL)


class FailureTest(PlanWorkGraphTestBase):
    def test_nonzero_exit_raises_planning_failed(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(returncode=1, stdout_obj=None, stderr="boom")
            with self.assertRaises(work_graph_plan.PlanningFailedError):
                work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")

    def test_is_error_true_raises_planning_failed(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(stdout_obj={"is_error": True})
            with self.assertRaises(work_graph_plan.PlanningFailedError):
                work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")

    def test_unparseable_result_raises_planning_failed(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(stdout_obj={"is_error": False, "result": "no json here"})
            with self.assertRaises(work_graph_plan.PlanningFailedError):
                work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")
        self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM work_graph_run").fetchone()[0], 0)


# E (2026-08-06): cross-Org-Agent dynamic planning path -- org_node != 'planner'
# gets _GENERIC_PROMPT_TEMPLATE instead of the fixed doc-chain prompt, and each
# node carries its own task_prompt instead of a fixed artifact_type.

_CODER_THREE_MODULE_MANIFEST = {
    "nodes": [
        {"node_id": "test-module-a", "depends_on": [], "task_prompt": "Write unit tests for module_a.py."},
        {"node_id": "test-module-b", "depends_on": [], "task_prompt": "Write unit tests for module_b.py."},
        {"node_id": "test-module-c", "depends_on": [], "task_prompt": "Write unit tests for module_c.py."},
    ],
    "edges": [],
}


class GenericPlanningPromptTest(PlanWorkGraphTestBase):
    def test_planner_gets_the_self_planning_catalog_template(self):
        # 2026-08-08 (FR-B4 "全部自由规划"): Planner no longer gets a fixed
        # "produce exactly BRD/PRD/SDD/Benchmark" instruction -- it gets the
        # node type catalog and decides itself which nodes this goal needs.
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")
        prompt = run.call_args.args[0][run.call_args.args[0].index("-p") + 1]
        self.assertIn("Available node types", prompt)
        self.assertIn("\"artifact_type\": \"benchmark\"", prompt)
        self.assertIn("on_failure", prompt)
        self.assertNotIn("design-dynamic-work-graph", prompt)
        self.assertNotIn("{goal_text}", prompt)  # the placeholder itself must have been substituted

    def test_coder_gets_generic_template_referencing_the_skill(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_CODER_THREE_MODULE_MANIFEST)}
            )
            work_graph_plan.plan_work_graph(self.conn, "coder", "task-1", "goal", "/tmp/proj-1")
        prompt = run.call_args.args[0][run.call_args.args[0].index("-p") + 1]
        self.assertIn("design-dynamic-work-graph", prompt)
        self.assertNotIn("BRD, PRD, SDD, Benchmark", prompt)
        self.assertIn("task_prompt", prompt)


class CoderThreeIndependentModulesScenarioTest(PlanWorkGraphTestBase):
    """E acceptance criterion 1: a Coder scenario that naturally splits into
    2-3 independent sub-tasks produces a graph_manifest via the new generic
    path (not the old single trigger_org_node call), correctly marking those
    nodes as mutually independent."""

    def test_three_independent_modules_produce_correctly_shaped_manifest(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_CODER_THREE_MODULE_MANIFEST)}
            )
            result = work_graph_plan.plan_work_graph(self.conn, "coder", "task-1", "goal", "/tmp/proj-1")

        manifest = result["graph_manifest"]
        self.assertEqual(len(manifest["nodes"]), 3)
        for node in manifest["nodes"]:
            self.assertEqual(node["depends_on"], [])
            self.assertIn("task_prompt", node)

        # E's whole point ties directly into G: independence is expressed
        # structurally (empty depends_on), so topological_layers() must put
        # all three in ONE concurrent layer, not three serial ones.
        layers = work_graph_orchestrate.topological_layers(manifest)
        self.assertEqual(len(layers), 1)
        self.assertEqual({n["node_id"] for n in layers[0]}, {"test-module-a", "test-module-b", "test-module-c"})


class GenericNodePromptBuilderTest(unittest.TestCase):
    def test_returns_the_node_own_task_prompt(self):
        node = {"node_id": "test-module-a", "task_prompt": "Write unit tests for module_a.py."}
        self.assertEqual(
            work_graph_plan.generic_node_prompt_builder(node, {}), "Write unit tests for module_a.py."
        )


class GenericPathEndToEndViaRunGraphTest(PlanWorkGraphTestBase):
    """E acceptance criterion (FR-5): the generated graph_manifest, driven by
    generic_node_prompt_builder, actually runs to completion through
    run_graph() -- not just a standalone parsing/shape check."""

    def test_three_independent_nodes_execute_and_return_outputs(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_CODER_THREE_MODULE_MANIFEST)}
            )
            plan_result = work_graph_plan.plan_work_graph(self.conn, "coder", "task-1", "goal", "/tmp/proj-1")

        # plan_work_graph() already inserted the work_graph_run row
        # (status='planning') -- run_graph() doesn't require any particular
        # status on it, so nothing more to set up here.
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph") as mock_execute:
            mock_execute.side_effect = lambda cwd, name, args, schema, timeout=600: {"content": "done: " + name}
            outputs = work_graph_orchestrate.run_graph(
                self.conn, plan_result["run_id"], "coder", "coder-implementation", "/tmp/proj-1",
                plan_result["graph_manifest"], work_graph_plan.generic_node_prompt_builder,
                {"type": "object", "required": ["node_id"]}, {"type": "object", "required": ["content"]},
            )

        self.assertEqual(set(outputs.keys()), {"test-module-a", "test-module-b", "test-module-c"})
        self.assertEqual(mock_execute.call_count, 3)


if __name__ == "__main__":
    unittest.main()