宜博 / lemongraph

"""D (2026-08-06) tests: work_graph_fork.py, mocked at the claude
subprocess boundary (CLAUDE.md Testing constraint) -- same pattern as
test_work_graph_orchestrate.py. Covers the requirements doc's acceptance
standards 1-3 (docs/done/work-graph-time-travel-debugging-requirements.md).

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

import json
import os
import tempfile
import unittest
from unittest.mock import patch

from src import schema, work_graph_execute, work_graph_fork


_BRD_PRD_SDD_BENCHMARK_MANIFEST = {
    "nodes": [
        {"node_id": "brd", "artifact_type": "brd", "depends_on": []},
        {"node_id": "prd", "artifact_type": "prd", "depends_on": ["brd"]},
        {"node_id": "sdd", "artifact_type": "sdd", "depends_on": ["prd"]},
        {"node_id": "benchmark", "artifact_type": "benchmark", "depends_on": ["sdd"]},
    ],
    "edges": [],
}


class ForkTestBase(unittest.TestCase):
    def setUp(self):
        self.home_tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.home_tmp.cleanup)
        self.old_graphlemon_home = os.environ.get("GRAPHLEMON_HOME")
        os.environ["GRAPHLEMON_HOME"] = self.home_tmp.name
        self.addCleanup(self._restore_graphlemon_home)

        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', 'validating', 'planner')"
        )
        self.conn.execute(
            "INSERT INTO skill (skill_id, owner_org_node, file_path, current_git_commit_sha, source, status) "
            "VALUES ('planner-doc-writing', 'planner', 'skills/planner-doc-writing/', 'unknown', 'new', 'active')"
        )
        self.conn.commit()

    def _restore_graphlemon_home(self):
        if self.old_graphlemon_home is None:
            os.environ.pop("GRAPHLEMON_HOME", None)
        else:
            os.environ["GRAPHLEMON_HOME"] = self.old_graphlemon_home

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

    def _seed_completed_run(self, run_id, org_node, graph_manifest, contents, status="done"):
        """Simulates an already-finished original run: a work_graph_run
        row plus one 'done' work_node_execution row per node, each holding
        contents[node_id] as its output_payload."""
        self.conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, graph_manifest, status) VALUES (?, 'task-1', ?, ?)",
            (run_id, json.dumps(graph_manifest), status),
        )
        node_exec_ids = {}
        for node in graph_manifest["nodes"]:
            node_id = node["node_id"]
            node_exec_id = run_id + "-" + node_id
            node_exec_ids[node_id] = node_exec_id
            self.conn.execute(
                "INSERT INTO work_node_execution (node_exec_id, run_id, node_id, subagent_type, skill_ref, "
                "input_schema_ref, output_schema_ref, input_payload, output_payload, status) "
                "VALUES (?, ?, ?, ?, 'planner-doc-writing', 'in.json', 'out.json', ?, ?, 'done')",
                (
                    node_exec_id, run_id, node_id, org_node,
                    json.dumps({"node_id": node_id}), json.dumps(contents[node_id]),
                ),
            )
        self.conn.commit()
        return node_exec_ids


class ForkFromPrdAcceptanceCriterion1Test(ForkTestBase):
    """Acceptance standard 1: fork from PRD with different content, SDD and
    Benchmark re-execute with the new PRD content, BRD is not re-executed."""

    def test_downstream_nodes_reexecute_upstream_node_reused(self):
        contents = {
            "brd": {"content": "original brd"}, "prd": {"content": "original prd"},
            "sdd": {"content": "original sdd"}, "benchmark": {"content": "original benchmark"},
        }
        node_exec_ids = self._seed_completed_run(
            "run-orig", "planner", _BRD_PRD_SDD_BENCHMARK_MANIFEST, contents
        )

        execute_calls = []

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            execute_calls.append(name)
            return {"content": "regenerated for " + name}

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            result = work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["prd"], {"content": "a completely different prd"}, "/tmp/proj-1",
            )

        # only SDD and Benchmark actually ran claude -p calls -- BRD and PRD did not.
        self.assertEqual(len(execute_calls), 2)

        outputs = result["outputs"]
        self.assertEqual(outputs["brd"], {"content": "original brd"})
        self.assertEqual(outputs["prd"], {"content": "a completely different prd"})
        self.assertNotEqual(outputs["sdd"], {"content": "original sdd"})
        self.assertNotEqual(outputs["benchmark"], {"content": "original benchmark"})

        new_run_id = result["run_id"]
        rows = {
            r[0]: (r[1], json.loads(r[2]))
            for r in self.conn.execute(
                "SELECT node_id, status, output_payload FROM work_node_execution WHERE run_id = ?", (new_run_id,)
            ).fetchall()
        }
        self.assertEqual(rows["brd"], ("done", {"content": "original brd"}))
        self.assertEqual(rows["prd"], ("done", {"content": "a completely different prd"}))
        self.assertEqual(rows["sdd"][0], "done")
        self.assertEqual(rows["benchmark"][0], "done")

        run_row = self.conn.execute(
            "SELECT status, parent_run_id, fork_of_node_exec_id FROM work_graph_run WHERE run_id = ?",
            (new_run_id,),
        ).fetchone()
        self.assertEqual(run_row, ("done", "run-orig", node_exec_ids["prd"]))


class ForkFromMixedShapeGraphTest(ForkTestBase):
    """Phase 2 §2.1.2 (2026-08-12) regression: a graph_manifest can now
    genuinely mix Planner-shaped nodes ('artifact_type') and Coder-shaped
    nodes ('task_prompt') side by side, once plan_work_graph(...,
    append_to_run_id=...) appends Coder's self-planned nodes onto
    Planner's original run. Forking from the artifact_type node must
    still correctly reconstruct a build_prompt for the DOWNSTREAM
    task_prompt node -- before this fix, _reconstruct_build_prompt
    required the WHOLE graph to be one uniform shape and would have
    raised ForkGraphShapeUnsupportedError here."""

    def test_downstream_generic_node_reexecutes_after_forking_upstream_planner_node(self):
        manifest = {
            "nodes": [
                {"node_id": "sdd", "artifact_type": "sdd", "depends_on": []},
                {
                    "node_id": "impl", "depends_on": ["sdd"], "task_prompt": "implement per the sdd",
                    "promotes_to_artifact": {"artifact_type": "test_report"},
                },
            ],
            "edges": [],
        }
        contents = {"sdd": {"content": "original sdd"}, "impl": {"content": "original impl"}}
        node_exec_ids = self._seed_completed_run("run-orig", "planner", manifest, contents)

        execute_calls = []

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            execute_calls.append(name)
            return {"content": "regenerated for " + name}

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            result = work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["sdd"], {"content": "a completely different sdd"}, "/tmp/proj-1",
            )

        # only "impl" (downstream of the fork point) genuinely re-ran.
        self.assertEqual(len(execute_calls), 1)
        outputs = result["outputs"]
        self.assertEqual(outputs["sdd"], {"content": "a completely different sdd"})
        self.assertNotEqual(outputs["impl"], {"content": "original impl"})


class OriginalRunUnaffectedTest(ForkTestBase):
    """Acceptance standard 2: the original run's records are completely
    untouched by a fork."""

    def test_original_run_and_nodes_unchanged_after_fork(self):
        contents = {
            "brd": {"content": "original brd"}, "prd": {"content": "original prd"},
            "sdd": {"content": "original sdd"}, "benchmark": {"content": "original benchmark"},
        }
        node_exec_ids = self._seed_completed_run(
            "run-orig", "planner", _BRD_PRD_SDD_BENCHMARK_MANIFEST, contents
        )
        before_run = self.conn.execute("SELECT * FROM work_graph_run WHERE run_id = 'run-orig'").fetchone()
        before_nodes = self.conn.execute(
            "SELECT * FROM work_node_execution WHERE run_id = 'run-orig' ORDER BY node_id"
        ).fetchall()

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "x"}):
            work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["prd"], {"content": "different"}, "/tmp/proj-1",
            )

        after_run = self.conn.execute("SELECT * FROM work_graph_run WHERE run_id = 'run-orig'").fetchone()
        after_nodes = self.conn.execute(
            "SELECT * FROM work_node_execution WHERE run_id = 'run-orig' ORDER BY node_id"
        ).fetchall()
        self.assertEqual(before_run, after_run)
        self.assertEqual(before_nodes, after_nodes)


class MultipleForksFromSameNodeTest(ForkTestBase):
    """Acceptance standard 3: the same historical node can be forked
    multiple times, producing independent, separately-traceable branches."""

    def test_two_forks_from_same_node_are_independent(self):
        contents = {
            "brd": {"content": "original brd"}, "prd": {"content": "original prd"},
            "sdd": {"content": "original sdd"}, "benchmark": {"content": "original benchmark"},
        }
        node_exec_ids = self._seed_completed_run(
            "run-orig", "planner", _BRD_PRD_SDD_BENCHMARK_MANIFEST, contents
        )

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "branch-a"}):
            result_a = work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["prd"], {"content": "prd variant A"}, "/tmp/proj-1",
            )
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "branch-b"}):
            result_b = work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["prd"], {"content": "prd variant B"}, "/tmp/proj-1",
            )

        self.assertNotEqual(result_a["run_id"], result_b["run_id"])
        self.assertEqual(result_a["outputs"]["prd"], {"content": "prd variant A"})
        self.assertEqual(result_b["outputs"]["prd"], {"content": "prd variant B"})
        self.assertEqual(result_a["outputs"]["sdd"], {"content": "branch-a"})
        self.assertEqual(result_b["outputs"]["sdd"], {"content": "branch-b"})

        for run_id in (result_a["run_id"], result_b["run_id"]):
            row = self.conn.execute(
                "SELECT parent_run_id, fork_of_node_exec_id FROM work_graph_run WHERE run_id = ?", (run_id,)
            ).fetchone()
            self.assertEqual(row, ("run-orig", node_exec_ids["prd"]))


class ForkValidationErrorsTest(ForkTestBase):
    def test_forking_a_non_done_node_raises(self):
        self.conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, graph_manifest, status) "
            "VALUES ('run-1', 'task-1', ?, 'running')",
            (json.dumps(_BRD_PRD_SDD_BENCHMARK_MANIFEST),),
        )
        self.conn.execute(
            "INSERT INTO work_node_execution (node_exec_id, run_id, node_id, subagent_type, skill_ref, "
            "input_schema_ref, output_schema_ref, input_payload, status) "
            "VALUES ('node-1', 'run-1', 'prd', 'planner', 'planner-doc-writing', 'in.json', 'out.json', '{}', 'running')"
        )
        self.conn.commit()
        with self.assertRaises(work_graph_fork.ForkSourceNotDoneError):
            work_graph_fork.fork_work_graph_run(self.conn, "node-1", {"content": "x"}, "/tmp/proj-1")

    def test_forking_from_a_non_terminal_parent_run_raises(self):
        contents = {"brd": {"content": "b"}, "prd": {"content": "p"}, "sdd": {"content": "s"}, "benchmark": {"content": "bm"}}
        node_exec_ids = self._seed_completed_run(
            "run-orig", "planner", _BRD_PRD_SDD_BENCHMARK_MANIFEST, contents, status="running",
        )
        with self.assertRaises(work_graph_fork.ForkParentRunNotTerminalError):
            work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["prd"], {"content": "x"}, "/tmp/proj-1",
            )

    def test_unsupported_graph_shape_raises(self):
        manifest = {"nodes": [{"node_id": "a", "depends_on": []}, {"node_id": "b", "depends_on": ["a"]}], "edges": []}
        node_exec_ids = self._seed_completed_run(
            "run-orig", "coder", manifest, {"a": {"content": "x"}, "b": {"content": "y"}},
        )
        with self.assertRaises(work_graph_fork.ForkGraphShapeUnsupportedError):
            work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["a"], {"content": "z"}, "/tmp/proj-1",
            )


class StartForkInBackgroundTest(unittest.TestCase):
    """Real threading (own connection per thread) -- needs a real file
    db_path, not ':memory:' (a fresh ':memory:' connection per thread
    would be a completely separate, empty database), same reasoning as
    test_task_execution.py's own StartTaskExecutionInBackgroundTest."""

    def setUp(self):
        self.home_tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.home_tmp.cleanup)
        self.old_graphlemon_home = os.environ.get("GRAPHLEMON_HOME")
        os.environ["GRAPHLEMON_HOME"] = self.home_tmp.name
        self.addCleanup(self._restore_graphlemon_home)

        self.tmp_db = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp_db.close()
        self.db_path = self.tmp_db.name
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

        conn = schema.connect(self.db_path)
        schema.create_schema(conn)
        conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
        )
        conn.execute(
            "INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) "
            "VALUES ('task-1', 'proj-1', 'g', 'goal', 'validating', 'planner')"
        )
        conn.execute(
            "INSERT INTO skill (skill_id, owner_org_node, file_path, current_git_commit_sha, source, status) "
            "VALUES ('planner-doc-writing', 'planner', 'skills/planner-doc-writing/', 'unknown', 'new', 'active')"
        )
        conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, graph_manifest, status) "
            "VALUES ('run-orig', 'task-1', ?, 'done')",
            (json.dumps(_BRD_PRD_SDD_BENCHMARK_MANIFEST),),
        )
        contents = {
            "brd": {"content": "b"}, "prd": {"content": "p"}, "sdd": {"content": "s"}, "benchmark": {"content": "bm"},
        }
        for node in _BRD_PRD_SDD_BENCHMARK_MANIFEST["nodes"]:
            node_id = node["node_id"]
            conn.execute(
                "INSERT INTO work_node_execution (node_exec_id, run_id, node_id, subagent_type, skill_ref, "
                "input_schema_ref, output_schema_ref, input_payload, output_payload, status) "
                "VALUES (?, 'run-orig', ?, 'planner', 'planner-doc-writing', 'in.json', 'out.json', ?, ?, 'done')",
                ("run-orig-" + node_id, node_id, json.dumps({"node_id": node_id}), json.dumps(contents[node_id])),
            )
        conn.commit()
        conn.close()

    def _restore_graphlemon_home(self):
        if self.old_graphlemon_home is None:
            os.environ.pop("GRAPHLEMON_HOME", None)
        else:
            os.environ["GRAPHLEMON_HOME"] = self.old_graphlemon_home

    def test_returns_run_id_immediately_and_completes_in_background(self):
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "regenerated"}):
            run_id, thread = work_graph_fork.start_fork_in_background(
                self.db_path, "run-orig-prd", {"content": "new prd"}, "/tmp/proj-1",
            )
            self.assertTrue(run_id)
            thread.join(timeout=10)

        conn = schema.connect(self.db_path)
        status = conn.execute("SELECT status FROM work_graph_run WHERE run_id = ?", (run_id,)).fetchone()[0]
        sdd_output = conn.execute(
            "SELECT output_payload FROM work_node_execution WHERE run_id = ? AND node_id = 'sdd'", (run_id,)
        ).fetchone()[0]
        conn.close()
        self.assertEqual(status, "done")
        self.assertEqual(json.loads(sdd_output), {"content": "regenerated"})


class ForkOfLeafNodeTest(ForkTestBase):
    """Forking the last node in a chain has nothing downstream -- no real
    execution needed, just the substituted output in a new run."""

    def test_forking_leaf_node_needs_no_execution(self):
        contents = {
            "brd": {"content": "original brd"}, "prd": {"content": "original prd"},
            "sdd": {"content": "original sdd"}, "benchmark": {"content": "original benchmark"},
        }
        node_exec_ids = self._seed_completed_run(
            "run-orig", "planner", _BRD_PRD_SDD_BENCHMARK_MANIFEST, contents
        )
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph") as mock_execute:
            mock_execute.side_effect = AssertionError("should not be called -- nothing downstream of benchmark")
            result = work_graph_fork.fork_work_graph_run(
                self.conn, node_exec_ids["benchmark"], {"content": "different benchmark"}, "/tmp/proj-1",
            )
        self.assertEqual(result["outputs"]["benchmark"], {"content": "different benchmark"})


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