宜博 / lemongraph

"""U-032 tests: the scheduler's own bookkeeping (loop-closing wiring + real
thread-level concurrency bookkeeping), mocked at the org_level_loop/
org_loop_route/run_goal_to_completion boundary -- this file is not about
re-proving e2e_pipeline/org_level_loop's own mechanics (covered in their own
test files), it's about proving finish_task_lifecycle applies the right
loop-closing step and run_scheduler_once actually runs goals concurrently
(not accidentally serially) while still collecting every result. Real
multi-Goal coverage (RUN_REAL_AGENT_E2E=1) is in src/test_scheduler_e2e.py.

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

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

from src import (
    bug_record, e2e_pipeline, org_level_loop, org_loop_route, paths, schema, scheduler, settings_config,
    work_graph_execute, work_node_run,
)


class FinishTaskLifecycleTestBase(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, loop_round, assigned_org_node) "
            "VALUES ('task-1', 'proj-1', 'goal', 'goal', 'validating', 0, 'validator')"
        )
        self.conn.commit()

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


class FinishTaskLifecyclePassTest(FinishTaskLifecycleTestBase):
    def test_pass_verdict_closes_task_done(self):
        run_result = {
            "task_id": "task-1", "doc_artifact_ids": {}, "code_artifact_ids": {},
            "verdict": {"verdict": "PASS", "raw_response": "VERDICT: PASS"},
        }
        result = scheduler.finish_task_lifecycle(self.conn, run_result, "/tmp/proj-1", "goal text")
        self.assertEqual(result["final_status"], "done")
        status = self.conn.execute("SELECT status FROM task WHERE task_id = 'task-1'").fetchone()[0]
        self.assertEqual(status, "done")


class FinishTaskLifecycleSelfTargetAuditTest(FinishTaskLifecycleTestBase):
    """U-035: self-target PASS closures additively record a core_engine
    audit-trail bug_record, without changing close_task_done's own effect."""

    def test_self_target_core_engine_change_records_audit_bug(self):
        run_result = {
            "task_id": "task-1", "doc_artifact_ids": {}, "code_artifact_ids": {"code_diff": "art-diff"},
            "verdict": {"verdict": "PASS", "raw_response": "VERDICT: PASS"},
        }
        with patch.object(paths, "is_self_target", return_value=True), \
             patch.object(e2e_pipeline, "_read_artifact_content", return_value=(
                 "diff --git a/src/scheduler.py b/src/scheduler.py\n"
                 "--- a/src/scheduler.py\n+++ b/src/scheduler.py\n"
             )):
            result = scheduler.finish_task_lifecycle(self.conn, run_result, "/tmp/proj-1", "goal text")

        self.assertEqual(result["final_status"], "done")
        status = self.conn.execute("SELECT status FROM task WHERE task_id = 'task-1'").fetchone()[0]
        self.assertEqual(status, "done")  # unaffected -- audit trail is additive, not a gate
        audit_row = self.conn.execute(
            "SELECT root_cause_category, status FROM bug_record WHERE task_id = 'task-1'"
        ).fetchone()
        self.assertEqual(audit_row, ("core_engine", "resolved"))

    def test_self_target_non_core_engine_change_records_no_audit_bug(self):
        run_result = {
            "task_id": "task-1", "doc_artifact_ids": {}, "code_artifact_ids": {"code_diff": "art-diff"},
            "verdict": {"verdict": "PASS", "raw_response": "VERDICT: PASS"},
        }
        with patch.object(paths, "is_self_target", return_value=True), \
             patch.object(e2e_pipeline, "_read_artifact_content", return_value=(
                 "diff --git a/frontend/app.js b/frontend/app.js\n"
                 "--- a/frontend/app.js\n+++ b/frontend/app.js\n"
             )):
            scheduler.finish_task_lifecycle(self.conn, run_result, "/tmp/proj-1", "goal text")

        count = self.conn.execute("SELECT COUNT(*) FROM bug_record WHERE task_id = 'task-1'").fetchone()[0]
        self.assertEqual(count, 0)

    def test_non_self_target_pass_records_no_audit_bug(self):
        run_result = {
            "task_id": "task-1", "doc_artifact_ids": {}, "code_artifact_ids": {"code_diff": "art-diff"},
            "verdict": {"verdict": "PASS", "raw_response": "VERDICT: PASS"},
        }
        # is_self_target NOT patched -- "/tmp/proj-1" is genuinely not this
        # repo's own checkout, so _read_artifact_content must never even be
        # attempted (would fail on a nonexistent artifact/file otherwise).
        scheduler.finish_task_lifecycle(self.conn, run_result, "/tmp/proj-1", "goal text")
        count = self.conn.execute("SELECT COUNT(*) FROM bug_record WHERE task_id = 'task-1'").fetchone()[0]
        self.assertEqual(count, 0)


class FinishTaskLifecycleFailTest(FinishTaskLifecycleTestBase):
    def test_fail_verdict_runs_fix_loop(self):
        run_result = {
            "task_id": "task-1",
            "doc_artifact_ids": {"benchmark": "artifact-b"},
            "code_artifact_ids": {},
            "verdict": {"verdict": "FAIL", "raw_response": "VERDICT: FAIL\nwall collision broken"},
        }
        with patch.object(scheduler.e2e_pipeline, "_read_artifact_content", return_value="benchmark content"), \
             patch.object(org_level_loop, "run_fix_loop", return_value={"final_status": "done"}) as mock_fix:
            result = scheduler.finish_task_lifecycle(self.conn, run_result, "/tmp/proj-1", "goal text")

        self.assertEqual(result["final_status"], "done")
        mock_fix.assert_called_once()
        args, kwargs = mock_fix.call_args
        self.assertEqual(args[0], self.conn)
        self.assertEqual(args[1], "task-1")
        self.assertEqual(args[2], "/tmp/proj-1")


class RunSchedulerOnceTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        conn = schema.connect(self.db_path)
        schema.create_schema(conn)
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def test_runs_all_goals_and_collects_results_in_order(self):
        def fake_run(db_path, target_folder_path, goal_text):
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-{}".format(i)) for i in range(4)]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=2)

        self.assertEqual([r["goal_text"] for r in results], [g[1] for g in goals])
        self.assertTrue(all("result" in r for r in results))

    def test_one_goal_raising_does_not_stop_the_others(self):
        def fake_run(db_path, target_folder_path, goal_text):
            if goal_text == "goal-bad":
                raise RuntimeError("boom")
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-ok-1"), ("/tmp/p", "goal-bad"), ("/tmp/p", "goal-ok-2")]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=3)

        self.assertIn("result", results[0])
        self.assertIn("error", results[1])
        self.assertEqual(results[1]["error"], "boom")
        self.assertIn("result", results[2])

    def test_max_workers_bounds_real_concurrency(self):
        current = [0]
        peak = [0]
        lock = threading.Lock()

        def fake_run(db_path, target_folder_path, goal_text):
            with lock:
                current[0] += 1
                peak[0] = max(peak[0], current[0])
            time.sleep(0.05)
            with lock:
                current[0] -= 1
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-{}".format(i)) for i in range(6)]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=2)

        self.assertEqual(len(results), 6)
        self.assertLessEqual(peak[0], 2)
        self.assertEqual(peak[0], 2, "expected real parallelism up to the max_workers bound, not accidental serial execution")

    def test_defaults_max_workers_to_org_concurrency_config(self):
        conn = schema.connect(self.db_path)
        settings_config.upsert_concurrency_config(conn, max_concurrent_goals=3, updated_by="test")
        conn.close()

        current = [0]
        peak = [0]
        lock = threading.Lock()

        def fake_run(db_path, target_folder_path, goal_text):
            with lock:
                current[0] += 1
                peak[0] = max(peak[0], current[0])
            time.sleep(0.05)
            with lock:
                current[0] -= 1
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-{}".format(i)) for i in range(6)]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            scheduler.run_scheduler_once(self.db_path, goals, max_workers=None)

        self.assertEqual(peak[0], 3)


class ConcurrentSqliteWritesStressTest(unittest.TestCase):
    """Not gated behind RUN_REAL_AGENT_E2E -- doesn't call the real Claude
    API, but DOES exercise the actual novel risk in this module's design
    (see module docstring's "Threading model" note): multiple threads, each
    holding its OWN sqlite connection, writing concurrently to the SAME db
    file. e2e_pipeline.run_end_to_end is faked (no real API calls) but
    finish_task_lifecycle's own close_task_done still runs for real against
    a real connection, so this is a genuine concurrent-write stress test,
    not a fully mocked one."""

    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        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.commit()
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def test_many_concurrent_goals_write_to_the_same_db_file_without_locking_errors(self):
        def fake_run_task_lifecycle(conn, task_id, target_folder_path, goal_text):
            # task_id already exists -- run_goal_to_completion creates it for
            # real via task_create.create_task before calling this; stand in
            # for org_node_lifecycle's real queued->...->validating march so
            # finish_task_lifecycle's close_task_done (requires 'validating') works.
            conn.execute("UPDATE task SET status = 'validating' WHERE task_id = ?", (task_id,))
            conn.commit()
            time.sleep(0.02)  # widen the contention window
            return {
                "task_id": task_id, "doc_artifact_ids": {}, "code_artifact_ids": {},
                "verdict": {"verdict": "PASS", "raw_response": "VERDICT: PASS"},
            }

        goals = [("/tmp/proj-1", "goal-{}".format(i)) for i in range(10)]
        with patch.object(scheduler.e2e_pipeline, "run_task_lifecycle", side_effect=fake_run_task_lifecycle):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=5)

        self.assertTrue(all("result" in r for r in results), results)

        conn = schema.connect(self.db_path)
        done_count = conn.execute("SELECT COUNT(*) FROM task WHERE status = 'done'").fetchone()[0]
        conn.close()
        self.assertEqual(done_count, 10)


class RoutableExceptionTestBase(unittest.TestCase):
    """REQ-06/07/14 failure: work-node-level failures (CorruptPersistedState
    /ClassifierBlocked/UsageLimitWaitTimeout) must actually reach
    task.status='blocked_pending_human' with the matching blocked_reason --
    found missing during the PH4-closure audit (2026-08-02): each exception
    was correctly raised and correctly marked its own work_node_execution
    'blocked', but nothing ever caught it to apply the task-level
    consequence 06's own acceptance text names."""

    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        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.commit()
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def _run_with_raising_lifecycle(self, exc):
        def fake_run_task_lifecycle(conn, task_id, target_folder_path, goal_text):
            raise exc

        with patch.object(scheduler.e2e_pipeline, "run_task_lifecycle", side_effect=fake_run_task_lifecycle):
            result = scheduler.run_goal_to_completion(self.db_path, "/tmp/proj-1", "goal text")
        return result

    def _assert_task_blocked(self, task_id, blocked_reason):
        conn = schema.connect(self.db_path)
        row = conn.execute("SELECT status, blocked_reason FROM task WHERE task_id = ?", (task_id,)).fetchone()
        event = conn.execute(
            "SELECT event_type FROM execution_event WHERE task_id = ? ORDER BY rowid DESC LIMIT 1", (task_id,)
        ).fetchone()
        conn.close()
        self.assertEqual(row, ("blocked_pending_human", blocked_reason))
        self.assertEqual(event[0], "task_status_changed")


class CorruptPersistedStateRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(work_node_run.CorruptPersistedStateError("boom"))
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "corrupt_persisted_state")
        self._assert_task_blocked(result["task_id"], "corrupt_persisted_state")

    def test_real_exception_message_is_persisted_for_diagnosis(self):
        # 2026-08-03: a blocked task used to carry only a short routing code
        # (e.g. 'execution_error') with the real failure reason nowhere to
        # be found -- found via a real console Start-button failure that
        # left no trace of *why*. _block_task now also stores the real
        # exception message in the execution_event payload.
        result = self._run_with_raising_lifecycle(work_node_run.CorruptPersistedStateError("disk full, boom"))
        conn = schema.connect(self.db_path)
        payload = conn.execute(
            "SELECT payload FROM execution_event WHERE task_id = ? ORDER BY rowid DESC LIMIT 1",
            (result["task_id"],),
        ).fetchone()[0]
        conn.close()
        self.assertEqual(json.loads(payload)["detail"], "disk full, boom")


class ClassifierBlockedRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(work_graph_execute.ClassifierBlockedError("boom"))
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "classifier_blocked")
        self._assert_task_blocked(result["task_id"], "classifier_blocked")


class UsageLimitWaitTimeoutRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(work_node_run.UsageLimitWaitTimeoutError("boom"))
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "usage_limit_wait_timeout")
        self._assert_task_blocked(result["task_id"], "usage_limit_wait_timeout")


class ConsistencyCheckFailedRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(
            scheduler.work_graph_orchestrate.ConsistencyCheckFailedError("boom")
        )
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "prd_consistency_check_failed")
        self._assert_task_blocked(result["task_id"], "prd_consistency_check_failed")


class UnroutedExceptionStillBlocksTaskTest(RoutableExceptionTestBase):
    def test_other_exceptions_block_the_task_instead_of_leaving_it_stuck(self):
        # 2026-08-04 regression: this used to propagate uncaught, leaving
        # the task neither running nor blocked_pending_human -- a real
        # work_graph_execute.WorkGraphExecuteError (API "Response stalled
        # mid-stream") crashed the calling process and left the task stuck
        # in limbo, nothing a human could resume from the console. Mirrors
        # task_execution._mark_blocked's same fallback for its own
        # (Start-button) entry point.
        result = self._run_with_raising_lifecycle(RuntimeError("something else entirely"))
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "execution_error")
        self._assert_task_blocked(result["task_id"], "execution_error")

    def test_real_exception_message_is_persisted_for_diagnosis(self):
        result = self._run_with_raising_lifecycle(RuntimeError("stalled mid-stream"))
        conn = schema.connect(self.db_path)
        payload = conn.execute(
            "SELECT payload FROM execution_event WHERE task_id = ? ORDER BY rowid DESC LIMIT 1",
            (result["task_id"],),
        ).fetchone()[0]
        conn.close()
        self.assertEqual(json.loads(payload)["detail"], "stalled mid-stream")


def _fix_loop_trigger(org_node, prompt, target_folder_path, timeout=300):
    """org_node_trigger.trigger_org_node fake for org_level_loop.run_fix_loop's
    3 perspective calls (all agree 'coder') + fixer + verify -- verify always
    fails, so route_org_loop keeps incrementing loop_round on retries."""
    if "CATEGORY:" in prompt:
        return "CATEGORY: coder\nDETAIL: because reasons"
    if org_node == "coder":
        return "Fixed it."
    if org_node == "validator":
        return "VERDICT: FAIL\nstill broken"
    raise AssertionError("unexpected call: " + org_node)


class SchedulerContinuesPastBlockedTaskTest(unittest.TestCase):
    """REQ-05 failure / REQ-09 boundary, against the REAL scheduler (U-032) --
    test_org_loop_route.py's own test for this literally says "no scheduler
    exists yet to test against directly" (stale as of U-032). One Goal's
    task is pre-seeded at loop_round=3 so its fix-loop attempt hits
    org_loop_route.LoopLimitReachedError for real (route_org_loop's own real
    blocked_pending_human write, not a mock) while the others complete via
    close_task_done -- proving the scheduler doesn't stop serving the rest
    of the batch."""

    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        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.commit()
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def test_one_task_hitting_loop_limit_does_not_block_the_others(self):
        def fake_run_task_lifecycle(conn, task_id, target_folder_path, goal_text):
            conn.execute("UPDATE task SET status = 'validating' WHERE task_id = ?", (task_id,))
            conn.commit()
            if goal_text == "goal-doomed":
                conn.execute("UPDATE task SET loop_round = 3 WHERE task_id = ?", (task_id,))
                conn.commit()
                verdict = {"verdict": "FAIL", "raw_response": "VERDICT: FAIL\nstill broken"}
            else:
                verdict = {"verdict": "PASS", "raw_response": "VERDICT: PASS"}
            return {
                "task_id": task_id, "doc_artifact_ids": {"benchmark": "artifact-b"},
                "code_artifact_ids": {}, "verdict": verdict,
            }

        goals = [("/tmp/proj-1", "goal-ok-1"), ("/tmp/proj-1", "goal-doomed"), ("/tmp/proj-1", "goal-ok-2")]
        with patch.object(scheduler.e2e_pipeline, "run_task_lifecycle", side_effect=fake_run_task_lifecycle), \
             patch.object(scheduler.e2e_pipeline, "_read_artifact_content", return_value="benchmark content"), \
             patch("src.org_node_trigger.trigger_org_node", side_effect=_fix_loop_trigger):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=3)

        by_goal = {r["goal_text"]: r for r in results}
        self.assertEqual(by_goal["goal-ok-1"]["result"]["final_status"], "done")
        self.assertEqual(by_goal["goal-ok-2"]["result"]["final_status"], "done")
        self.assertIn("error", by_goal["goal-doomed"])

        conn = schema.connect(self.db_path)
        doomed_status = conn.execute(
            "SELECT status, blocked_reason FROM task WHERE goal_text = 'goal-doomed'"
        ).fetchone()
        conn.close()
        self.assertEqual(doomed_status, ("blocked_pending_human", "loop_limit_reached"))


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