宜博 / lemongraph

"""U-031 tests: work_graph_orchestrate.py -- the graph_manifest -> per-node
execution bridge closed during the REQ-01~14 regression audit. Real sqlite,
run_fn mocked at the claude subprocess boundary (execute_work_graph is
tested for real elsewhere, in test_work_graph_execute.py / the 2026-08-01
spike -- this module's own job is orchestration, not re-proving that).

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

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

from src import execution_log, paths, schema, work_graph_execute, work_graph_orchestrate as wgo, work_node_run


class TopologicalOrderTest(unittest.TestCase):
    def test_linear_chain_via_depends_on(self):
        manifest = {
            "nodes": [
                {"node_id": "brd", "depends_on": []},
                {"node_id": "prd", "depends_on": ["brd"]},
                {"node_id": "sdd", "depends_on": ["prd"]},
            ],
            "edges": [],
        }
        self.assertEqual([n["node_id"] for n in wgo.topological_order(manifest)], ["brd", "prd", "sdd"])

    def test_linear_chain_via_edges_only(self):
        # real LLM output might populate edges but not each node's own
        # depends_on -- must still order correctly
        manifest = {
            "nodes": [{"node_id": "a"}, {"node_id": "b"}, {"node_id": "c"}],
            "edges": [{"from": "a", "to": "b"}, {"from": "b", "to": "c"}],
        }
        self.assertEqual([n["node_id"] for n in wgo.topological_order(manifest)], ["a", "b", "c"])

    def test_diamond_fan_out_fan_in(self):
        manifest = {
            "nodes": [
                {"node_id": "n1", "depends_on": []},
                {"node_id": "n2", "depends_on": ["n1"]},
                {"node_id": "n3", "depends_on": ["n1"]},
                {"node_id": "n4", "depends_on": ["n2", "n3"]},
            ],
            "edges": [],
        }
        order = [n["node_id"] for n in wgo.topological_order(manifest)]
        self.assertEqual(order[0], "n1")
        self.assertEqual(order[-1], "n4")
        self.assertEqual(set(order[1:3]), {"n2", "n3"})

    def test_cycle_raises(self):
        manifest = {
            "nodes": [{"node_id": "x", "depends_on": ["y"]}, {"node_id": "y", "depends_on": ["x"]}],
            "edges": [],
        }
        with self.assertRaises(ValueError):
            wgo.topological_order(manifest)

    def test_independent_nodes_ordered_deterministically(self):
        manifest = {"nodes": [{"node_id": "b"}, {"node_id": "a"}], "edges": []}
        self.assertEqual([n["node_id"] for n in wgo.topological_order(manifest)], ["a", "b"])


class TopologicalLayersTest(unittest.TestCase):
    """G: topological_layers() -- the layer-grouped algorithm run_graph()'s
    concurrent scheduler is built on. topological_order() is now a thin
    flattening wrapper over this (tested above, unchanged behavior)."""

    def test_diamond_produces_three_layers(self):
        manifest = {
            "nodes": [
                {"node_id": "n1", "depends_on": []},
                {"node_id": "n2", "depends_on": ["n1"]},
                {"node_id": "n3", "depends_on": ["n1"]},
                {"node_id": "n4", "depends_on": ["n2", "n3"]},
            ],
            "edges": [],
        }
        layers = wgo.topological_layers(manifest)
        self.assertEqual(
            [[n["node_id"] for n in layer] for layer in layers],
            [["n1"], ["n2", "n3"], ["n4"]],
        )

    def test_linear_chain_is_all_single_node_layers(self):
        manifest = {
            "nodes": [
                {"node_id": "brd", "depends_on": []},
                {"node_id": "prd", "depends_on": ["brd"]},
                {"node_id": "sdd", "depends_on": ["prd"]},
            ],
            "edges": [],
        }
        layers = wgo.topological_layers(manifest)
        self.assertEqual(
            [[n["node_id"] for n in layer] for layer in layers],
            [["brd"], ["prd"], ["sdd"]],
        )

    def test_independent_nodes_are_a_single_layer(self):
        manifest = {"nodes": [{"node_id": "b"}, {"node_id": "a"}], "edges": []}
        layers = wgo.topological_layers(manifest)
        self.assertEqual(len(layers), 1)
        self.assertEqual([n["node_id"] for n in layers[0]], ["a", "b"])

    def test_cycle_raises(self):
        manifest = {
            "nodes": [{"node_id": "x", "depends_on": ["y"]}, {"node_id": "y", "depends_on": ["x"]}],
            "edges": [],
        }
        with self.assertRaises(ValueError):
            wgo.topological_layers(manifest)

    def test_topological_order_matches_flattened_layers(self):
        manifest = {
            "nodes": [
                {"node_id": "n1", "depends_on": []},
                {"node_id": "n2", "depends_on": ["n1"]},
                {"node_id": "n3", "depends_on": ["n1"]},
            ],
            "edges": [],
        }
        layers = wgo.topological_layers(manifest)
        flat = [n["node_id"] for layer in layers for n in layer]
        self.assertEqual([n["node_id"] for n in wgo.topological_order(manifest)], flat)


class RunGraphTestBase(unittest.TestCase):
    def setUp(self):
        # DU-01: run_fn now really calls execution_log.append_log(), which
        # writes to paths.log_file_path() (default ~/.graphlemon/logs/...) --
        # isolate that under a throwaway tempdir so real (non-mocked)
        # run_work_node/run_fn calls in this test module never touch the
        # developer's actual home directory, same pattern as test_paths.py.
        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)

        # G: check_same_thread=False -- run_graph()'s concurrent layers call
        # into this conn (via _LockedConn) from ThreadPoolExecutor worker
        # threads, not just the thread that created the connection.
        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', 'goal', 'coding', 'planner')"
        )
        self.conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) "
            "VALUES ('run-1', 'task-1', 'planner', '{}', 'running')"
        )
        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()


_OUT_SCHEMA = {"type": "object", "required": ["content"]}
_IN_SCHEMA = {"type": "object", "required": ["node_id"]}


class RunGraphHappyPathTest(RunGraphTestBase):
    def test_runs_nodes_in_order_and_returns_outputs(self):
        manifest = {
            "nodes": [
                {"node_id": "brd", "artifact_type": "brd", "depends_on": []},
                {"node_id": "prd", "artifact_type": "prd", "depends_on": ["brd"]},
            ],
            "edges": [],
        }
        calls = []

        def build_prompt(node, outputs):
            calls.append((node["node_id"], dict(outputs)))
            return "prompt for " + node["node_id"]

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph") as mock_execute:
            mock_execute.side_effect = [{"content": "brd text"}, {"content": "prd text"}]
            outputs = wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                build_prompt, _IN_SCHEMA, _OUT_SCHEMA,
            )

        self.assertEqual(outputs, {"brd": {"content": "brd text"}, "prd": {"content": "prd text"}})
        self.assertEqual(calls[0], ("brd", {}))
        self.assertEqual(calls[1], ("prd", {"brd": {"content": "brd text"}}))

    def test_creates_work_node_execution_rows_and_placeholder_skill(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "x"}):
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
            )
        row = self.conn.execute(
            "SELECT status, subagent_type, skill_ref FROM work_node_execution WHERE run_id = 'run-1' AND node_id = 'brd'"
        ).fetchone()
        self.assertEqual(row, ("done", "planner", "planner-doc-writing"))
        skill_row = self.conn.execute("SELECT owner_org_node FROM skill WHERE skill_id = 'planner-doc-writing'").fetchone()
        self.assertEqual(skill_row[0], "planner")

    def test_resumes_without_rerunning_done_nodes(self):
        manifest = {
            "nodes": [
                {"node_id": "brd", "depends_on": []},
                {"node_id": "prd", "depends_on": ["brd"]},
            ],
            "edges": [],
        }
        calls = []

        def build_prompt(node, outputs):
            return "p-" + node["node_id"]

        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: calls.append(name) or {"content": "c-" + name}
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                build_prompt, _IN_SCHEMA, _OUT_SCHEMA,
            )
        self.assertEqual(len(calls), 2)

        # simulate a restart: run_graph called again for the same run_id
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph") as mock_execute2:
            mock_execute2.side_effect = AssertionError("should not be called -- both nodes are already done")
            outputs = wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                build_prompt, _IN_SCHEMA, _OUT_SCHEMA,
            )
        self.assertEqual(set(outputs.keys()), {"brd", "prd"})


class RunGraphUsageLimitTest(RunGraphTestBase):
    def test_usage_limit_reached_translates_to_quota_wait(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        attempts = {"n": 0}

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            attempts["n"] += 1
            if attempts["n"] == 1:
                raise work_graph_execute.UsageLimitReachedError("hit your session limit", "session")
            return {"content": "recovered"}

        sleeps = []
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            outputs = wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                sleep=lambda s: sleeps.append(s), now=lambda: 0.0,
            )
        self.assertEqual(outputs["brd"], {"content": "recovered"})
        self.assertEqual(sleeps, [5 * 60])


class RunGraphClassifierBlockedTest(RunGraphTestBase):
    def test_classifier_blocked_marks_node_blocked_and_propagates(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(
                 work_graph_execute, "execute_work_graph",
                 side_effect=work_graph_execute.ClassifierBlockedError("blocked by classifier"),
             ):
            with self.assertRaises(work_graph_execute.ClassifierBlockedError):
                wgo.run_graph(
                    self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                    lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                )
        node_exec_id = self.conn.execute(
            "SELECT node_exec_id FROM work_node_execution WHERE run_id = 'run-1' AND node_id = 'brd'"
        ).fetchone()[0]
        status = self.conn.execute(
            "SELECT status FROM work_node_execution WHERE node_exec_id = ?", (node_exec_id,)
        ).fetchone()[0]
        self.assertEqual(status, "blocked")


class VerifyFnTest(RunGraphTestBase):
    """REQ-02 boundary: verify_fn intercepts a bad node output before it
    propagates downstream, via one regeneration attempt."""

    def test_verify_fn_passing_immediately_does_not_regenerate(self):
        manifest = {"nodes": [{"node_id": "prd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "consistent prd"}) as mock_execute:
            outputs = wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                verify_fn=lambda node, output: (True, None),
            )
        self.assertEqual(outputs["prd"], {"content": "consistent prd"})
        self.assertEqual(mock_execute.call_count, 1)
        run_status = self.conn.execute("SELECT status FROM work_graph_run WHERE run_id = 'run-1'").fetchone()[0]
        self.assertEqual(run_status, "running")  # untouched -- verify_fn never failed

    def test_verify_fn_failing_once_regenerates_and_succeeds(self):
        manifest = {"nodes": [{"node_id": "prd", "depends_on": []}], "edges": []}
        contents = iter(["contradictory prd", "fixed prd"])
        checked = []

        def verify_fn(node, output):
            checked.append(output["content"])
            return output["content"] == "fixed prd", ["scenario A contradicts scenario B"]

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=lambda *a, **k: {"content": next(contents)}):
            outputs = wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA, verify_fn=verify_fn,
            )
        self.assertEqual(outputs["prd"], {"content": "fixed prd"})
        self.assertEqual(checked, ["contradictory prd", "fixed prd"])
        run_status = self.conn.execute("SELECT status FROM work_graph_run WHERE run_id = 'run-1'").fetchone()[0]
        self.assertEqual(run_status, "verifying")
        reflect_row = self.conn.execute(
            "SELECT status FROM work_node_execution WHERE run_id = 'run-1' AND node_id = 'prd-reflect'"
        ).fetchone()
        self.assertEqual(reflect_row[0], "done")

    def test_verify_fn_failing_twice_raises_and_blocks_run(self):
        manifest = {"nodes": [{"node_id": "prd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "still contradictory"}):
            with self.assertRaises(wgo.ConsistencyCheckFailedError):
                wgo.run_graph(
                    self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                    lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                    verify_fn=lambda node, output: (False, ["still broken"]),
                )
        run_status = self.conn.execute("SELECT status FROM work_graph_run WHERE run_id = 'run-1'").fetchone()[0]
        self.assertEqual(run_status, "blocked")


class SelfTargetForcesWorktreeTest(RunGraphTestBase):
    """U-035 (05 §5.8): self-target unconditionally forces needs_worktree=True,
    regardless of what the caller passed. run_work_node itself is mocked here
    (not just execute_work_graph) so this only asserts what needs_worktree
    value run_graph computes and passes down -- it doesn't exercise real git
    worktree commands against this actual checkout."""

    def test_self_target_forces_worktree_even_when_caller_passed_false(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        self.conn.execute(
            "UPDATE project SET target_folder_path = ? WHERE project_id = 'proj-1'", (paths.repo_root(),)
        )
        self.conn.commit()
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_node_run, "run_work_node", return_value={"content": "x"}) as mock_run:
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", paths.repo_root(), manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA, needs_worktree=False,
            )
        self.assertTrue(mock_run.call_args.kwargs["needs_worktree"])

    def test_non_self_target_respects_caller_supplied_false(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_node_run, "run_work_node", return_value={"content": "x"}) as mock_run:
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA, needs_worktree=False,
            )
        self.assertFalse(mock_run.call_args.kwargs["needs_worktree"])

    def test_non_self_target_still_respects_caller_supplied_true(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_node_run, "run_work_node", return_value={"content": "x"}) as mock_run:
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA, needs_worktree=True,
            )
        self.assertTrue(mock_run.call_args.kwargs["needs_worktree"])


class RunFnLogsExecutionEventsTest(RunGraphTestBase):
    """DU-01 (bug A1/A2 root cause): `_make_run_fn()`'s `run_fn` closure must
    actually call `execution_log.append_log()` around the real
    `execute_work_graph()` call, so `execution_log_file`'s backing file on
    disk genuinely has content -- not just a DB row tracking its lifecycle."""

    def _node_exec_id(self, node_id="brd"):
        return self.conn.execute(
            "SELECT node_exec_id FROM work_node_execution WHERE run_id = 'run-1' AND node_id = ?", (node_id,)
        ).fetchone()[0]

    def _log_lines(self, node_exec_id):
        content = execution_log.get_log(self.conn, node_exec_id)["content"]
        return [json.loads(line) for line in content.splitlines() if line]

    def test_success_writes_node_start_then_node_completed(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "brd text"}):
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
            )

        events = self._log_lines(self._node_exec_id())
        self.assertEqual([e["event"] for e in events], ["node_start", "node_completed"])
        self.assertEqual(events[1]["output"], {"content": "brd text"})

        # also exercise the real HTTP-facing read path end to end (not just
        # the module function) -- an empty file was exactly bug A1.
        row = self.conn.execute(
            "SELECT size_bytes FROM execution_log_file WHERE node_exec_id = ?", (self._node_exec_id(),)
        ).fetchone()
        self.assertGreater(row[0], 0)

    def test_classifier_blocked_writes_node_failed_and_still_propagates(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(
                 work_graph_execute, "execute_work_graph",
                 side_effect=work_graph_execute.ClassifierBlockedError("blocked by classifier"),
             ):
            with self.assertRaises(work_graph_execute.ClassifierBlockedError):
                wgo.run_graph(
                    self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                    lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                )

        events = self._log_lines(self._node_exec_id())
        self.assertEqual([e["event"] for e in events], ["node_start", "node_failed"])
        self.assertIn("blocked by classifier", events[1]["error"])

    def test_usage_limit_writes_quota_limit_then_retries(self):
        manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []}
        attempts = {"n": 0}

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            attempts["n"] += 1
            if attempts["n"] == 1:
                raise work_graph_execute.UsageLimitReachedError("hit your session limit", "session")
            return {"content": "recovered"}

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            outputs = wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                sleep=lambda s: None, now=lambda: 0.0,
            )
        self.assertEqual(outputs["brd"], {"content": "recovered"})

        events = self._log_lines(self._node_exec_id())
        self.assertEqual(
            [e["event"] for e in events], ["node_start", "quota_limit", "node_start", "node_completed"]
        )
        self.assertEqual(events[1]["quota_limit_type"], "session")


class RunGraphConcurrencyTest(RunGraphTestBase):
    """G, design doc §10 acceptance criteria 1-2: same-layer nodes actually
    run concurrently (overlapping wall-clock intervals), and the next layer
    never starts until every node in the previous layer has finished."""

    def test_same_layer_nodes_run_concurrently(self):
        manifest = {
            "nodes": [{"node_id": "a", "depends_on": []}, {"node_id": "b", "depends_on": []}],
            "edges": [],
        }
        events = []
        events_lock = threading.Lock()

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            start = time.monotonic()
            time.sleep(0.05)
            end = time.monotonic()
            with events_lock:
                events.append((start, end))
            return {"content": name}

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p-" + node["node_id"], _IN_SCHEMA, _OUT_SCHEMA,
            )

        self.assertEqual(len(events), 2)
        (start_a, end_a), (start_b, end_b) = events
        self.assertLess(start_a, end_b)
        self.assertLess(start_b, end_a)

    def test_next_layer_waits_for_entire_previous_layer(self):
        manifest = {
            "nodes": [
                {"node_id": "a", "depends_on": []},
                {"node_id": "b", "depends_on": []},
                {"node_id": "c", "depends_on": ["a", "b"]},
            ],
            "edges": [],
        }
        events = []
        events_lock = threading.Lock()

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            start = time.monotonic()
            time.sleep(0.03)
            end = time.monotonic()
            with events_lock:
                events.append((start, end))
            return {"content": "x"}

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
            )

        self.assertEqual(len(events), 3)
        layer1_ends = [end for _, end in events[:2]]
        layer2_start, _ = events[2]
        self.assertGreaterEqual(layer2_start, max(layer1_ends))


class RunGraphLayerErrorTest(RunGraphTestBase):
    """G, design doc §10 acceptance criterion (FR-5): one node failing in a
    concurrent layer must not cancel an already-running sibling, and all
    failures in the layer are aggregated into one WorkGraphLayerError."""

    def test_one_failing_node_does_not_cancel_sibling_and_aggregates_errors(self):
        manifest = {
            "nodes": [{"node_id": "a", "depends_on": []}, {"node_id": "b", "depends_on": []}],
            "edges": [],
        }
        calls = []

        def fake_run_work_node(conn, node_exec_id, input_payload, input_schema, output_schema, run_fn, **kwargs):
            calls.append(input_payload["node_id"])
            if input_payload["node_id"] == "a":
                raise RuntimeError("boom on a")
            return {"content": "b done"}

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_node_run, "run_work_node", side_effect=fake_run_work_node):
            with self.assertRaises(wgo.WorkGraphLayerError) as ctx:
                wgo.run_graph(
                    self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                    lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                )

        self.assertEqual(sorted(calls), ["a", "b"])  # b was actually attempted, not cancelled
        error_node_ids = [node_id for node_id, exc in ctx.exception.errors]
        self.assertEqual(error_node_ids, ["a"])


class RunGraphConcurrentResumeTest(RunGraphTestBase):
    """G, design doc §10 acceptance criterion 3: breakpoint-recovery still
    works when the interrupted run had a concurrent (2+ node) layer -- a
    second run_graph() call for the same run_id must not re-run nodes the
    first call already finished."""

    def test_concurrent_layer_resume_reruns_nothing(self):
        manifest = {
            "nodes": [{"node_id": "a", "depends_on": []}, {"node_id": "b", "depends_on": []}],
            "edges": [],
        }
        calls = []

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            calls.append(name)
            return {"content": "c-" + name}

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
            )
        self.assertEqual(len(calls), 2)

        with patch.object(work_graph_execute, "write_workflow_script"), \
             patch.object(work_graph_execute, "execute_work_graph") as mock_execute2:
            mock_execute2.side_effect = AssertionError("should not be called -- both nodes already done")
            outputs = wgo.run_graph(
                self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
            )
        self.assertEqual(set(outputs.keys()), {"a", "b"})


class RunGraphQuotaCoexistsWithNormalCompletionTest(RunGraphTestBase):
    """G, design doc §10 acceptance criterion 6: a node stuck on a quota-wait
    timeout in a concurrent layer must not lose its sibling's real, already-
    persisted output -- the sibling's row is 'done' with its real payload
    even though the overall layer raises."""

    def test_quota_timeout_on_one_node_does_not_erase_siblings_persisted_output(self):
        manifest = {
            "nodes": [{"node_id": "a", "depends_on": []}, {"node_id": "b", "depends_on": []}],
            "edges": [],
        }
        prompt_by_script = {}
        prompt_lock = threading.Lock()

        def fake_write_script(cwd, script_name, script_content):
            with prompt_lock:
                prompt_by_script[script_name] = script_content

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            content = prompt_by_script[name]
            if '"p-a"' in content:
                raise work_graph_execute.UsageLimitReachedError("hit your session limit", "session")
            return {"content": "b done"}

        now_calls = {"n": 0}
        now_lock = threading.Lock()

        def fake_now():
            # first call records quota_wait_started_at as "0"; second call
            # (computing elapsed) reports far past the session timeout, so
            # this times out immediately instead of really sleeping.
            with now_lock:
                now_calls["n"] += 1
                return 0.0 if now_calls["n"] == 1 else 999999.0

        with patch.object(work_graph_execute, "write_workflow_script", side_effect=fake_write_script), \
             patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
            with self.assertRaises(wgo.WorkGraphLayerError) as ctx:
                wgo.run_graph(
                    self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                    lambda node, outputs: "p-" + node["node_id"], _IN_SCHEMA, _OUT_SCHEMA,
                    sleep=lambda s: None, now=fake_now,
                )

        error_node_ids = [node_id for node_id, exc in ctx.exception.errors]
        self.assertEqual(error_node_ids, ["a"])
        self.assertIsInstance(dict(ctx.exception.errors)["a"], work_node_run.UsageLimitWaitTimeoutError)

        b_row = self.conn.execute(
            "SELECT status, output_payload FROM work_node_execution WHERE run_id = 'run-1' AND node_id = 'b'"
        ).fetchone()
        self.assertEqual(b_row[0], "done")
        self.assertEqual(json.loads(b_row[1]), {"content": "b done"})


class RunGraphRealSqliteConcurrencyTest(unittest.TestCase):
    """G, design doc §10 acceptance criterion 4: a real (non-:memory:,
    non-mocked-conn) sqlite database under genuine concurrent write
    pressure from a multi-node layer must never raise 'database is
    locked' -- run several rounds since races like this often only
    reproduce probabilistically."""

    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.db_tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.db_tmp.cleanup)

    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_multiple_rounds_of_concurrent_layer_writes_never_raise_locked(self):
        manifest = {
            "nodes": [{"node_id": n, "depends_on": []} for n in ("a", "b", "c", "d")],
            "edges": [],
        }

        def fake_execute(cwd, name, args, output_schema, timeout=600):
            time.sleep(0.005)
            return {"content": name}

        for round_n in range(5):
            db_path = os.path.join(self.db_tmp.name, "test_{}.db".format(round_n))
            conn = schema.connect(db_path, check_same_thread=False)
            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', 'goal', 'goal', 'coding', 'planner')"
            )
            conn.execute(
                "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) "
                "VALUES ('run-1', 'task-1', 'planner', '{}', 'running')"
            )
            conn.commit()

            try:
                with patch.object(work_graph_execute, "write_workflow_script"), \
                     patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute):
                    outputs = wgo.run_graph(
                        conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest,
                        lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA,
                    )
                self.assertEqual(set(outputs.keys()), {"a", "b", "c", "d"})
            finally:
                conn.close()


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