codex-app-server/cloud/tests/test_worker_concurrency.py

"""Worker 并发单测。

只用注入的假 runner（假 CLI），不跑真实模型、不联网。
capacity>1 时同一 worker 内多个任务必须真正并行跑，而不是
被单任务字段或共享连接拖回串行；capacity=1 时行为与旧串行
完全一致；最终状态写回必须是 CAS，取消/回收后的任务不能
被后台线程盖回 completed。
"""

import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from types import SimpleNamespace

try:
    from cloud.server import init_db
    from cloud.worker import Worker
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.server import init_db
    from cloud.worker import Worker


def _insert_task(db_path, task_id, prompt="do work", project_id="p1"):
    now = time.time()
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT INTO tasks(id, project_id, prompt, status,"
            " idempotency_key, depends_on, priority, created_by,"
            " created_at, updated_at, last_heartbeat)"
            " VALUES(?, ?, ?, 'queued', NULL, '[]', 0, 'tester', ?, ?, NULL)",
            (task_id, project_id, prompt, now, now),
        )
        conn.commit()
    finally:
        conn.close()


def _status(db_path, task_id):
    conn = sqlite3.connect(db_path)
    try:
        row = conn.execute(
            "SELECT status FROM tasks WHERE id = ?",
            (task_id,)).fetchone()
    finally:
        conn.close()
    return row[0] if row else None


def _host_json(name, status, text):
    return SimpleNamespace(
        returncode=0,
        stdout=json.dumps({"task": name, "status": status, "text": text}),
        stderr="")


class WorkerConcurrencyTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        init_db(self.db_path)
        self.env = {"XCM_API_KEY": "test-placeholder"}

    def tearDown(self):
        self.tmp.cleanup()

    def _worker(self, runner=None, **kwargs):
        kwargs.setdefault("env", self.env)
        kwargs.setdefault("retry_backoff_seconds", 0)
        kwargs.setdefault("max_retries", 0)
        return Worker(self.db_path, "/fake/host.py",
                      runner=runner, **kwargs)

    def test_capacity_2_runs_two_tasks_concurrently(self):
        # 两个任务各睡 0.5 秒：并行时它们必须在同一时间段内
        # 都在跑（用 Barrier 证明重叠，而不是只看总时长，避免
        # 慢 CI 上的时间断言抖动）。串行时第一个 runner 会在
        # Barrier 上等满 5 秒超时，测试会因此失败——这正是要
        # 证明的：capacity=2 不是串行。
        barrier = threading.Barrier(2, timeout=5)

        def fake(cmd):
            barrier.wait()
            time.sleep(0.1)
            return _host_json("x", "completed", "done")

        _insert_task(self.db_path, "A")
        _insert_task(self.db_path, "B")
        worker = self._worker(runner=fake, capacity=2)
        results = worker.run_all()
        self.assertEqual(len(results), 2)
        self.assertEqual({r["status"] for r in results}, {"completed"})
        self.assertEqual(_status(self.db_path, "A"), "completed")
        self.assertEqual(_status(self.db_path, "B"), "completed")

    def test_capacity_1_stays_serial(self):
        # capacity=1 时 run_all 仍是旧串行循环：两个任务都完成，
        # 且第二个任务的认领发生在第一个完成之后（靠任务状态
        # 变化顺序断言，不靠时间）。
        order = []
        lock = threading.Lock()

        def fake(cmd):
            with lock:
                order.append("start")
            time.sleep(0.2)
            with lock:
                order.append("end")
            return _host_json("x", "completed", "done")

        _insert_task(self.db_path, "A")
        _insert_task(self.db_path, "B")
        worker = self._worker(runner=fake, capacity=1)
        results = worker.run_all()
        self.assertEqual(len(results), 2)
        self.assertEqual(order, ["start", "end", "start", "end"])

    def test_pool_max_tasks_respected(self):
        # 并行模式下 max_tasks=1 只跑一个任务：提交数受
        # max_tasks 约束，不能因为线程池先提交了两个就多跑。
        def fake(cmd):
            time.sleep(0.1)
            return _host_json("x", "completed", "done")

        for tid in ("A", "B", "C"):
            _insert_task(self.db_path, tid)
        worker = self._worker(runner=fake, capacity=2)
        results = worker.run_all(max_tasks=1)
        self.assertEqual(len(results), 1)
        done = [tid for tid in ("A", "B", "C")
                if _status(self.db_path, tid) == "completed"]
        self.assertEqual(len(done), 1)

    def test_wait_for_running(self):
        # _wait_for_running：有任务在跑时超时返回 False，
        # 任务清掉后返回 True（run_loop 优雅排空靠它）。
        worker = self._worker(capacity=2)
        worker._running_set("A", container="codex-A")
        self.assertFalse(worker._wait_for_running(timeout=0.3))

        def clearer():
            time.sleep(0.2)
            worker._running_clear("A")

        threading.Thread(target=clearer, daemon=True).start()
        self.assertTrue(worker._wait_for_running(timeout=5))

    def test_final_write_does_not_overwrite_interrupted(self):
        # 任务在跑时被外部标成 interrupted（取消/回收），随后
        # runner 返回 completed，最终写回绝不能把它盖回 completed。
        started = threading.Event()

        def fake(cmd):
            started.set()
            time.sleep(0.5)
            return _host_json("x", "completed", "done")

        _insert_task(self.db_path, "A")
        worker = self._worker(runner=fake, capacity=1)
        outcome = {}

        def run():
            outcome["result"] = worker.run_one()

        t = threading.Thread(target=run, daemon=True)
        t.start()
        self.assertTrue(started.wait(timeout=5))
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "UPDATE tasks SET status = 'interrupted' WHERE id = 'A'")
            conn.commit()
        finally:
            conn.close()
        t.join(timeout=10)
        self.assertEqual(outcome["result"]["status"], "interrupted")
        self.assertEqual(_status(self.db_path, "A"), "interrupted")


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