"""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()