宜博 / lemongraph

"""U-013 tests: c_work_node_run's contract (schema validation, retry state
machine, quota-wait, execution_log_file lifecycle) with a fake `run_fn`, plus
real git worktree isolation + merge -- REQ-03's boundary case (two independent
nodes, separate worktrees, merge without conflict) needs no LLM call at all,
so it's tested for real here, not mocked.

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

import json
import os
import subprocess
import tempfile
import unittest

from src import schema, work_node_run

_IN_SCHEMA = {"type": "object", "required": ["goal"]}
_OUT_SCHEMA = {"type": "object", "required": ["value"]}


def _init_git_repo(path):
    subprocess.run(["git", "init"], cwd=path, check=True, capture_output=True)
    with open(os.path.join(path, "README.md"), "w") as f:
        f.write("seed\n")
    # matches graphlemon's own .gitignore (2026-08-03) -- without this,
    # ensure_worktree_committed's plain `git add -A` on repo_path picks up
    # a still-present worktree checkout as a tracked gitlink-like entry.
    with open(os.path.join(path, ".gitignore"), "w") as f:
        f.write(".graphlemon-worktrees/\n")
    subprocess.run(["git", "add", "README.md", ".gitignore"], cwd=path, check=True, capture_output=True)
    subprocess.run(
        ["git"] + work_node_run._GIT_IDENTITY + ["commit", "-m", "seed"],
        cwd=path, check=True, capture_output=True,
    )


class WorkNodeRunTestBase(unittest.TestCase):
    def setUp(self):
        self.conn = schema.connect(":memory:")
        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', 'coder')"
        )
        self.conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) "
            "VALUES ('run-1', 'task-1', 'coder', '{}', 'running')"
        )
        self.conn.execute(
            "INSERT INTO skill (skill_id, owner_org_node, file_path, current_git_commit_sha, source, status) "
            "VALUES ('skill-1', 'coder', 'skills/skill-1/', 'sha', 'new', 'active')"
        )
        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', 'n1', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'pending')"
        )
        self.conn.commit()

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


class InputValidationTest(WorkNodeRunTestBase):
    def test_invalid_input_rejected_before_execution_starts(self):
        with self.assertRaises(work_node_run.InvalidInputError):
            work_node_run.run_work_node(
                self.conn, "node-1", {}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=lambda cwd, model=None: {"value": 1}
            )
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "pending")
        count = self.conn.execute("SELECT COUNT(*) FROM execution_log_file").fetchone()[0]
        self.assertEqual(count, 0)


class HappyPathTest(WorkNodeRunTestBase):
    def test_valid_output_marks_node_done(self):
        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=lambda cwd, model=None: {"value": 42}
        )
        self.assertEqual(output, {"value": 42})
        row = self.conn.execute(
            "SELECT status, output_payload, retry_count FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row[0], "done")
        self.assertEqual(json.loads(row[1]), {"value": 42})
        self.assertEqual(row[2], 0)

    def test_execution_log_file_created_active_then_closed(self):
        seen_status_during_run = []

        def run_fn(cwd, model=None):
            seen_status_during_run.append(
                self.conn.execute("SELECT status FROM execution_log_file").fetchone()[0]
            )
            return {"value": 1}

        work_node_run.run_work_node(self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn)
        self.assertEqual(seen_status_during_run, ["active"])
        final_status = self.conn.execute("SELECT status FROM execution_log_file").fetchone()[0]
        self.assertEqual(final_status, "closed")

    def test_log_file_closed_even_when_run_fn_raises(self):
        def boom(cwd, model=None):
            raise RuntimeError("node blew up")

        with self.assertRaises(RuntimeError):
            work_node_run.run_work_node(self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=boom)
        final_status = self.conn.execute("SELECT status FROM execution_log_file").fetchone()[0]
        self.assertEqual(final_status, "closed")

    def test_unretryable_run_fn_failure_marks_node_blocked_not_stuck_running(self):
        # any exception from run_fn other than a recognized usage-limit
        # signal (e.g. CLASSIFIER_BLOCKED surfacing from
        # work_graph_execute.execute_work_graph, 03-C §3.5 point 6) must
        # not leave the row stuck at status='running' with no record of why
        class SomeUnretryableError(Exception):
            pass

        def boom(cwd, model=None):
            raise SomeUnretryableError("classifier blocked or similar")

        with self.assertRaises(SomeUnretryableError):
            work_node_run.run_work_node(self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=boom)
        status = self.conn.execute(
            "SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()[0]
        self.assertEqual(status, "blocked")


class RetryOnSchemaFailureTest(WorkNodeRunTestBase):
    def test_retries_then_succeeds(self):
        attempts = {"n": 0}

        def run_fn(cwd, model=None):
            attempts["n"] += 1
            if attempts["n"] < 3:
                return {"wrong_key": True}
            return {"value": "ok"}

        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn, max_retries=2
        )
        self.assertEqual(output, {"value": "ok"})
        row = self.conn.execute(
            "SELECT status, retry_count FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row, ("done", 2))

    def test_schema_failure_passes_through_reflecting_before_retrying(self):
        # REQ-03 failure: "status 应为 reflecting->retrying" -- previously
        # jumped straight to 'retrying', 'reflecting' was never written.
        # execution_log_file/work_node_execution have no separate history
        # table, so the only way to observe the transient write is sqlite's
        # own trace callback (every SQL statement it actually executes).
        executed_sql = []
        self.conn.set_trace_callback(executed_sql.append)

        attempts = {"n": 0}

        def run_fn(cwd, model=None):
            attempts["n"] += 1
            if attempts["n"] == 1:
                return {"wrong_key": True}
            return {"value": "ok"}

        try:
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn, max_retries=2
            )
        finally:
            self.conn.set_trace_callback(None)

        reflecting_writes = [sql for sql in executed_sql if "status = 'reflecting'" in sql]
        self.assertEqual(len(reflecting_writes), 1)

    def test_exceeding_max_retries_blocks_the_node(self):
        with self.assertRaises(work_node_run.SchemaValidationFailedError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
                run_fn=lambda cwd, model=None: {"wrong_key": True}, max_retries=2,
            )
        row = self.conn.execute(
            "SELECT status, retry_count FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row, ("blocked", 3))


class UsageLimitWaitTest(WorkNodeRunTestBase):
    def setUp(self):
        super().setUp()
        self.sleeps = []
        self.clock = [0.0]

    def _sleep(self, seconds):
        self.sleeps.append(seconds)
        self.clock[0] += seconds

    def _now(self):
        return self.clock[0]

    def test_session_limit_then_recovers(self):
        attempts = {"n": 0}

        def run_fn(cwd, model=None):
            attempts["n"] += 1
            if attempts["n"] == 1:
                work_node_run.raise_if_usage_limited("Error: You've hit your session limit")
            return {"value": "recovered"}

        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now,
        )
        self.assertEqual(output, {"value": "recovered"})
        self.assertEqual(self.sleeps, [5 * 60])

        row = self.conn.execute(
            "SELECT quota_limit_type FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row[0], "session")
        event = self.conn.execute(
            "SELECT event_type FROM execution_event WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(event[0], "quota_wait_started")

    def test_opus_quota_retries_immediately_on_downgraded_model_no_backoff(self):
        # REQ-14 boundary: "不走退避轮询，直接降级到非 Opus 模型重试" -- an actual
        # retry with a different model, not just detect-and-fail.
        calls = []

        def run_fn(cwd, model=None):
            calls.append(model)
            if len(calls) == 1:
                work_node_run.raise_if_usage_limited("You've hit your Opus limit")
            return {"value": "recovered on " + str(model)}

        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now,
        )
        self.assertEqual(calls, [None, "claude-sonnet-5"])
        self.assertEqual(output, {"value": "recovered on claude-sonnet-5"})
        self.assertEqual(self.sleeps, [])  # no backoff wait for opus
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "done")

    def test_opus_quota_custom_downgrade_model_is_used(self):
        calls = []

        def run_fn(cwd, model=None):
            calls.append(model)
            if len(calls) == 1:
                work_node_run.raise_if_usage_limited("You've hit your Opus limit")
            return {"value": "ok"}

        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now, opus_downgrade_model="claude-fable-5",
        )
        self.assertEqual(calls, [None, "claude-fable-5"])

    def test_opus_quota_persisting_after_downgrade_raises_timeout(self):
        # a non-opus model shouldn't hit an opus-specific quota again -- if it
        # somehow does, this doesn't loop forever chasing it
        def run_fn(cwd, model=None):
            work_node_run.raise_if_usage_limited("You've hit your Opus limit")

        with self.assertRaises(work_node_run.UsageLimitWaitTimeoutError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
                sleep=self._sleep, now=self._now,
            )
        self.assertEqual(self.sleeps, [])

    def test_session_timeout_after_8_hours_blocks_node(self):
        def run_fn(cwd, model=None):
            work_node_run.raise_if_usage_limited("You've hit your session limit")

        with self.assertRaises(work_node_run.UsageLimitWaitTimeoutError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
                sleep=self._sleep, now=self._now,
            )
        self.assertGreaterEqual(self.clock[0], 8 * 3600)
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "blocked")

    def test_weekly_limit_then_recovers(self):
        # previously untested quota_limit_type -- session/opus were covered,
        # weekly wasn't, even though the detection regex and timeout table
        # both already supported it
        attempts = {"n": 0}

        def run_fn(cwd, model=None):
            attempts["n"] += 1
            if attempts["n"] == 1:
                work_node_run.raise_if_usage_limited("Error: You've hit your weekly limit")
            return {"value": "recovered"}

        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
            sleep=self._sleep, now=self._now,
        )
        quota_type = self.conn.execute(
            "SELECT quota_limit_type FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()[0]
        self.assertEqual(quota_type, "weekly")

    def test_weekly_timeout_after_9_days_blocks_node(self):
        def run_fn(cwd, model=None):
            work_node_run.raise_if_usage_limited("You've hit your weekly limit")

        with self.assertRaises(work_node_run.UsageLimitWaitTimeoutError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
                sleep=self._sleep, now=self._now,
            )
        self.assertGreaterEqual(self.clock[0], 9 * 24 * 3600)
        status = self.conn.execute("SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'").fetchone()[0]
        self.assertEqual(status, "blocked")


class RealGitWorktreeIsolationTest(WorkNodeRunTestBase):
    """REQ-03 boundary: two independent nodes, each in its own git worktree,
    merged back with no conflict. Real git, no LLM call needed."""

    def setUp(self):
        super().setUp()
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        _init_git_repo(self.tmp.name)
        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-2', 'run-1', 'n2', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'pending')"
        )
        self.conn.commit()

    def _write_file_run_fn(self, filename, content):
        def run_fn(cwd, model=None):
            with open(os.path.join(cwd, filename), "w") as f:
                f.write(content)
            subprocess.run(["git", "add", filename], cwd=cwd, check=True, capture_output=True)
            subprocess.run(
                ["git"] + work_node_run._GIT_IDENTITY + ["commit", "-m", "add " + filename],
                cwd=cwd, check=True, capture_output=True,
            )
            return {"value": filename}

        return run_fn

    def test_two_independent_nodes_isolated_and_merged_without_conflict(self):
        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=self._write_file_run_fn("module_a.txt", "module a content"),
            needs_worktree=True, repo_path=self.tmp.name,
        )
        work_node_run.run_work_node(
            self.conn, "node-2", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=self._write_file_run_fn("module_b.txt", "module b content"),
            needs_worktree=True, repo_path=self.tmp.name,
        )

        self.assertTrue(os.path.isfile(os.path.join(self.tmp.name, "module_a.txt")))
        self.assertTrue(os.path.isfile(os.path.join(self.tmp.name, "module_b.txt")))

        row1 = self.conn.execute(
            "SELECT git_worktree_ref FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()[0]
        row2 = self.conn.execute(
            "SELECT git_worktree_ref FROM work_node_execution WHERE node_exec_id = 'node-2'"
        ).fetchone()[0]
        self.assertNotEqual(row1, row2)
        self.assertIn("node-1", row1)
        self.assertIn("node-2", row2)

        log = subprocess.run(
            ["git", "log", "--oneline"], cwd=self.tmp.name, check=True, capture_output=True, text=True
        )
        self.assertIn("module_a.txt", log.stdout)
        self.assertIn("module_b.txt", log.stdout)

    def _write_file_never_commit_run_fn(self, filename, content):
        """Unlike _write_file_run_fn, deliberately edits and leaves the
        worktree DIRTY -- reproducing a real Coder Work Node, which is
        never instructed to `git commit` its own edits (found via a real
        self-hosting run, 2026-08-02: a real, verified Coder edit was
        silently discarded because nothing committed it before merge)."""
        def run_fn(cwd, model=None):
            with open(os.path.join(cwd, filename), "w") as f:
                f.write(content)
            return {"value": filename}

        return run_fn

    def test_uncommitted_worktree_changes_are_auto_committed_before_merge(self):
        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=self._write_file_never_commit_run_fn("uncommitted.txt", "real edit, never committed"),
            needs_worktree=True, repo_path=self.tmp.name,
        )

        merged_path = os.path.join(self.tmp.name, "uncommitted.txt")
        self.assertTrue(os.path.isfile(merged_path), "the Coder's edit must survive merge even if it never committed")
        with open(merged_path) as f:
            self.assertEqual(f.read(), "real edit, never committed")

        status = subprocess.run(
            ["git", "status", "--porcelain"], cwd=self.tmp.name, check=True, capture_output=True, text=True
        ).stdout
        self.assertEqual(status.strip(), "", "main checkout must be clean after merge, not carrying a stray dirty file")

    def test_clean_worktree_with_no_changes_is_still_a_valid_merge(self):
        # A node that legitimately makes no changes (e.g. a pure
        # investigation, see the real self-hosting run this fix responds
        # to) must not error just because there's nothing to auto-commit.
        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=lambda cwd, model=None: {"value": "investigated, no changes needed"},
            needs_worktree=True, repo_path=self.tmp.name,
        )
        status = subprocess.run(
            ["git", "status", "--porcelain"], cwd=self.tmp.name, check=True, capture_output=True, text=True
        ).stdout
        self.assertEqual(status.strip(), "")

    def test_drift_on_main_checkout_is_auto_committed_before_merge(self):
        # 2026-08-03 regression: git worktrees deliberately share project
        # identity with their main checkout -- a real self-hosting run
        # showed `claude`'s own memory writes can land on the MAIN
        # checkout (repo_path) instead of the worktree, regardless of
        # cwd. If left uncommitted, the NEXT node's own worktree merge
        # fails with "local changes would be overwritten by merge". This
        # run_fn simulates that: it writes a file directly into repo_path
        # (not cwd, i.e. not the worktree) to reproduce main-checkout
        # drift caused by something outside the node's own isolated work.
        def run_fn(cwd, model=None):
            with open(os.path.join(self.tmp.name, "main-checkout-drift.txt"), "w") as f:
                f.write("landed on the main checkout, not the worktree")
            return {"value": "done"}

        work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=run_fn, needs_worktree=True, repo_path=self.tmp.name,
        )

        status = subprocess.run(
            ["git", "status", "--porcelain"], cwd=self.tmp.name, check=True, capture_output=True, text=True
        ).stdout
        self.assertEqual(status.strip(), "", "main checkout drift must be auto-committed, not left dirty")
        self.assertTrue(os.path.isfile(os.path.join(self.tmp.name, "main-checkout-drift.txt")))

        # a SECOND node's own worktree merge must not fail because of it
        work_node_run.run_work_node(
            self.conn, "node-2", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=self._write_file_run_fn("module_b.txt", "module b content"),
            needs_worktree=True, repo_path=self.tmp.name,
        )
        self.assertTrue(os.path.isfile(os.path.join(self.tmp.name, "module_b.txt")))


class BreakpointRecoveryTest(WorkNodeRunTestBase):
    """U-021 / REQ-06: run_work_node called again for a node_exec_id that was
    already mid-flight when a process restart happened."""

    def test_done_node_returns_stored_output_without_rerunning(self):
        self.conn.execute(
            "UPDATE work_node_execution SET status = 'done', output_payload = ? WHERE node_exec_id = 'node-1'",
            (json.dumps({"value": "already done"}),),
        )
        self.conn.commit()
        calls = []
        output = work_node_run.run_work_node(
            self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
            run_fn=lambda cwd, model=None: calls.append(1) or {"value": "should not run"},
        )
        self.assertEqual(output, {"value": "already done"})
        self.assertEqual(calls, [])

    def test_resuming_retrying_node_continues_persisted_retry_count(self):
        # node-1 was already at retry_count=2 (not 0) before the restart --
        # one more failure must immediately exceed max_retries=2, proving
        # the count resumed rather than restarting from zero.
        self.conn.execute(
            "UPDATE work_node_execution SET status = 'retrying', retry_count = 2 WHERE node_exec_id = 'node-1'"
        )
        self.conn.commit()
        calls = []

        def run_fn(cwd, model=None):
            calls.append(1)
            return {"wrong_key": True}

        with self.assertRaises(work_node_run.SchemaValidationFailedError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn, max_retries=2,
            )
        self.assertEqual(len(calls), 1)
        row = self.conn.execute(
            "SELECT status, retry_count FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row, ("blocked", 3))

    def test_resuming_waiting_for_quota_node_uses_persisted_timestamp_not_a_fresh_clock(self):
        # if elapsed-time tracking restarted at 0 on resume instead of using
        # the persisted quota_wait_started_at, this would never time out.
        original_ts = "2026-07-01 00:00:00"
        epoch_orig = work_node_run._parse_sqlite_timestamp(original_ts)
        self.conn.execute(
            "UPDATE work_node_execution SET status = 'waiting_for_quota', "
            "quota_wait_started_at = ?, quota_limit_type = 'session' WHERE node_exec_id = 'node-1'",
            (original_ts,),
        )
        self.conn.commit()

        def run_fn(cwd, model=None):
            work_node_run.raise_if_usage_limited("You've hit your session limit")

        with self.assertRaises(work_node_run.UsageLimitWaitTimeoutError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA, run_fn=run_fn,
                sleep=lambda s: None, now=lambda: epoch_orig + 8 * 3600 + 1,
            )
        status = self.conn.execute(
            "SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()[0]
        self.assertEqual(status, "blocked")

    def test_nonexistent_node_exec_id_raises_corrupt_state(self):
        with self.assertRaises(work_node_run.CorruptPersistedStateError):
            work_node_run.run_work_node(
                self.conn, "no-such-node", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
                run_fn=lambda cwd, model=None: {"value": 1},
            )

    def test_done_node_with_null_output_payload_raises_corrupt_state(self):
        self.conn.execute(
            "UPDATE work_node_execution SET status = 'done', output_payload = NULL WHERE node_exec_id = 'node-1'"
        )
        self.conn.commit()
        with self.assertRaises(work_node_run.CorruptPersistedStateError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
                run_fn=lambda cwd, model=None: {"value": 1},
            )

    def test_waiting_for_quota_node_with_null_timestamp_raises_corrupt_state(self):
        self.conn.execute(
            "UPDATE work_node_execution SET status = 'waiting_for_quota', quota_wait_started_at = NULL "
            "WHERE node_exec_id = 'node-1'"
        )
        self.conn.commit()
        with self.assertRaises(work_node_run.CorruptPersistedStateError):
            work_node_run.run_work_node(
                self.conn, "node-1", {"goal": "x"}, _IN_SCHEMA, _OUT_SCHEMA,
                run_fn=lambda cwd, model=None: {"value": 1},
            )


class ResumeWorkGraphRunTest(WorkNodeRunTestBase):
    """U-021 / REQ-06: resume_work_graph_run's own resume-plan + corrupt-state
    detection, independent of run_work_node."""

    def test_returns_only_non_done_node_exec_ids(self):
        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 ('node-2', 'run-1', 'n2', 'coder', 'skill-1', 'in.json', 'out.json', '{}', ?, 'done')",
            (json.dumps({"value": 1}),),
        )
        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-3', 'run-1', 'n3', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'retrying')"
        )
        self.conn.commit()
        resumable = work_node_run.resume_work_graph_run(self.conn, "run-1")
        self.assertEqual(sorted(resumable), ["node-1", "node-3"])

    def test_nonexistent_run_raises_corrupt_state(self):
        with self.assertRaises(work_node_run.CorruptPersistedStateError):
            work_node_run.resume_work_graph_run(self.conn, "no-such-run")

    def test_invalid_run_status_raises_corrupt_state(self):
        # unreachable via a normal write (work_graph_run.status has its own
        # CHECK constraint) -- ignore_check_constraints simulates the kind of
        # externally-corrupted row REQ-06 must still detect.
        self.conn.execute("PRAGMA ignore_check_constraints = ON")
        self.conn.execute("UPDATE work_graph_run SET status = 'bogus_status' WHERE run_id = 'run-1'")
        self.conn.commit()
        with self.assertRaises(work_node_run.CorruptPersistedStateError):
            work_node_run.resume_work_graph_run(self.conn, "run-1")

    def test_done_node_missing_output_payload_raises_corrupt_state(self):
        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-2', 'run-1', 'n2', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'done')"
        )
        self.conn.commit()
        with self.assertRaises(work_node_run.CorruptPersistedStateError):
            work_node_run.resume_work_graph_run(self.conn, "run-1")

    def test_waiting_for_quota_node_missing_timestamp_raises_corrupt_state(self):
        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-2', 'run-1', 'n2', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'waiting_for_quota')"
        )
        self.conn.commit()
        with self.assertRaises(work_node_run.CorruptPersistedStateError):
            work_node_run.resume_work_graph_run(self.conn, "run-1")


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