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

"""多 worker 池派发 + 通用 CloudExecutor + 自动 PR 单测。

只用标准库，不跑真实模型回合、不联网、不登录任何服务；
git/gh 一律用假 runner 注入，不依赖本机是否装 gh、是否有
远端仓库；CloudExecutor 的 endpoint 一律用假 http_client。
"""

import http.client
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.auth import AuthStore
    from cloud.pr import PrSkipped, create_task_pr
    from cloud.sandbox import CloudExecutor, DockerExecutor
    from cloud.server import create_server, init_db
    from cloud.worker import Worker
    from cloud.workers import (
        list_workers,
        register_worker,
    )
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.pr import PrSkipped, create_task_pr
    from cloud.sandbox import CloudExecutor, DockerExecutor
    from cloud.server import create_server, init_db
    from cloud.worker import Worker
    from cloud.workers import (
        list_workers,
        register_worker,
    )


def _result(returncode=0, stdout="", stderr=""):
    return SimpleNamespace(returncode=returncode, stdout=stdout,
                           stderr=stderr)


class _ServerCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        self.auth = AuthStore(self.db_path)
        self.owner_id, self.owner_key = self.auth.create_user("owner")
        self.viewer_id, self.viewer_key = self.auth.create_user("viewer")
        self.server = create_server(self.db_path, "127.0.0.1", 0)
        self.port = self.server.server_address[1]
        self.thread = threading.Thread(
            target=self.server.serve_forever, daemon=True)
        self.thread.start()
        status, self.project = self._req(
            "POST", "/projects", {"name": "p1", "auto_pr": 1},
            self.owner_key)
        self.assertEqual(status, 201, self.project)
        self.auth.set_member(
            self.project["id"], self.owner_id, self.viewer_id, "viewer")
        self.worktrees = os.path.join(self.tmp.name, "worktrees")
        os.makedirs(self.worktrees, exist_ok=True)
        self.server.host_worktrees_dir = self.worktrees

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)
        try:
            self.auth.conn.close()
        except Exception:
            pass
        self.tmp.cleanup()

    def _req(self, method, path, body=None, key=None, headers=None):
        conn = http.client.HTTPConnection(
            "127.0.0.1", self.port, timeout=10)
        hdrs = dict(headers or {})
        payload = None
        if body is not None:
            payload = json.dumps(body).encode("utf-8")
            hdrs.setdefault("Content-Type", "application/json")
        if key is not None:
            hdrs["X-API-Key"] = key
        try:
            conn.request(method, path, body=payload, headers=hdrs)
            resp = conn.getresponse()
            data = resp.read()
        finally:
            conn.close()
        try:
            parsed = json.loads(data.decode("utf-8")) if data else None
        except ValueError:
            parsed = None
        return resp.status, parsed

    def _insert_task(self, task_id, status, project_id=None):
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "INSERT INTO tasks(id, project_id, prompt, status,"
                " depends_on, priority, created_by, created_at,"
                " updated_at, last_heartbeat) VALUES(?, ?, ?, ?, '[]',"
                " 0, 'system', 0, 0, NULL)",
                (task_id, project_id or self.project["id"],
                 "do it", status))
            conn.commit()
        finally:
            conn.close()

    def _task_row(self, task_id):
        conn = sqlite3.connect(self.db_path)
        try:
            row = conn.execute(
                "SELECT status, pr_url, pr_number FROM tasks"
                " WHERE id = ?", (task_id,)).fetchone()
        finally:
            conn.close()
        return row

    def _audit_actions(self):
        conn = sqlite3.connect(self.db_path)
        try:
            return [r[0] for r in conn.execute(
                "SELECT action FROM audit_log").fetchall()]
        finally:
            conn.close()

    def _git_runner(self, push_fail=False, remote_fail=False):
        def run(cmd, cwd):
            if "remote" in cmd:
                if remote_fail:
                    return _result(1, "", "no remote")
                return _result(0, "https://example.com/repo.git\n")
            if "branch" in cmd:
                # 假 worktree 正好在自己的 host/<任务名> 分支上：
                # worktree 目录名就是任务 id（见 _worktree）。
                return _result(0, "host/%s\n" % os.path.basename(cwd or ""))
            if "add" in cmd or "commit" in cmd:
                return _result(0, "")
            if "push" in cmd:
                if push_fail:
                    return _result(1, "", "push rejected")
                return _result(0, "")
            return _result(0, "")
        return run

    def _gh_runner(self, authed=True, pr_url=None):
        def run(cmd, cwd):
            if cmd[:2] == ["gh", "auth"]:
                return _result(0 if authed else 1, "", "")
            if cmd[:2] == ["gh", "pr"]:
                return _result(0, (pr_url or
                                   "https://example.com/repo/pull/42\n"))
            return _result(0, "")
        return run

    def _worktree(self, task_id):
        path = os.path.join(self.worktrees, task_id)
        os.makedirs(path, exist_ok=True)
        return path


class WorkerRegistryTestCase(_ServerCase):
    def test_heartbeat_requires_worker_id(self):
        status, body = self._req(
            "POST", "/workers/heartbeat", {}, self.owner_key)
        self.assertEqual(status, 400, body)

    def test_heartbeat_registers_and_lists_alive(self):
        status, body = self._req(
            "POST", "/workers/heartbeat",
            {"worker_id": "w-docker-1", "executor": "docker",
             "capacity": 2},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertTrue(body["alive"])
        self.assertEqual(body["executor"], "docker")
        # P1：完整 workers 清单只限全局管理员，本测试里 owner 即管理员。
        old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
        os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
        try:
            status, body = self._req("GET", "/workers", None, self.owner_key)
        finally:
            if old_admins is None:
                os.environ.pop("CODEX_GLOBAL_ADMINS", None)
            else:
                os.environ["CODEX_GLOBAL_ADMINS"] = old_admins
        self.assertEqual(status, 200, body)
        self.assertEqual(len(body["workers"]), 1)
        self.assertTrue(body["workers"][0]["alive"])

    def test_heartbeat_preserves_executor_when_not_given(self):
        self._req("POST", "/workers/heartbeat",
                  {"worker_id": "w1", "executor": "docker"}, self.owner_key)
        status, body = self._req(
            "POST", "/workers/heartbeat", {"worker_id": "w1"},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["executor"], "docker")

    def test_stale_worker_lists_not_alive(self):
        register_worker(self.db_path, "w-old", executor="local")
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "UPDATE workers SET last_seen = ? WHERE worker_id = ?",
                (time.time() - 3600, "w-old"))
            conn.commit()
        finally:
            conn.close()
        workers = list_workers(self.db_path)
        self.assertEqual(len(workers), 1)
        self.assertFalse(workers[0]["alive"])

    def test_heartbeat_without_key_is_unauthorized(self):
        status, body = self._req(
            "POST", "/workers/heartbeat", {"worker_id": "w1"}, None)
        self.assertEqual(status, 401, body)


class TaskExecutorFilterTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        init_db(os.path.join(self.tmp.name, "control.db"))
        self.db_path = os.path.join(self.tmp.name, "control.db")
        conn = sqlite3.connect(self.db_path)
        try:
            for tid, exc in (("t-local", "local"), ("t-none", None),
                             ("t-docker", "docker")):
                conn.execute(
                    "INSERT INTO tasks(id, project_id, prompt, status,"
                    " depends_on, priority, created_by, created_at,"
                    " updated_at, last_heartbeat, executor)"
                    " VALUES(?, 'proj', 'do it', 'queued', '[]', 0,"
                    " 'system', 0, 0, NULL, ?)", (tid, exc))
            conn.commit()
        finally:
            conn.close()

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

    def _docker_worker(self, capacity=None):
        return Worker(
            self.db_path, "host.py", config_path=None, repo=None,
            executor=DockerExecutor(runner=lambda cmd: _result()),
            env={"XCM_API_KEY": "test-placeholder"}, env_key=None,
            task_timeout=1800, max_retries=0, retry_backoff_seconds=0,
            capacity=capacity)

    def test_claim_rejects_other_pool_task(self):
        worker = self._docker_worker()
        self.assertFalse(worker._claim("t-local", executor="docker"))
        conn = sqlite3.connect(self.db_path)
        try:
            status = conn.execute(
                "SELECT status FROM tasks WHERE id = 't-local'"
            ).fetchone()[0]
        finally:
            conn.close()
        self.assertEqual(status, "queued")

    def test_claim_accepts_unassigned_and_same_pool(self):
        # 本测试连续认领两个任务, 测的是执行器过滤不是上限;
        # caps 默认上限 1 会挡第二单, 故显式给 capacity=2。
        worker = self._docker_worker(capacity=2)
        self.assertTrue(worker._claim("t-none", executor="docker"))
        self.assertTrue(worker._claim("t-docker", executor="docker"))

    def test_cas_race_only_one_worker_wins(self):
        w1 = self._docker_worker()
        w2 = self._docker_worker()
        self.assertTrue(w1._claim("t-none", executor="docker"))
        self.assertFalse(w2._claim("t-none", executor="docker"))

    def test_pick_next_skips_other_pool(self):
        # 只剩 local 池任务时，docker worker 无就绪任务。
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute("DELETE FROM tasks WHERE id != 't-local'")
            conn.commit()
        finally:
            conn.close()
        worker = self._docker_worker()
        self.assertIsNone(worker.pick_next())

    def test_run_one_skips_other_pool_without_model(self):
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute("DELETE FROM tasks WHERE id != 't-local'")
            conn.commit()
        finally:
            conn.close()
        worker = self._docker_worker()
        result = worker.run_one()
        self.assertFalse(result["dispatched"])
        self.assertEqual(result["reason"], "no_ready_task")


class CloudExecutorTestCase(unittest.TestCase):
    def test_unconfigured_fails_without_fallback(self):
        ex = CloudExecutor(endpoint=None, env={})
        with self.assertRaises(RuntimeError) as ctx:
            ex.execute({"id": "t1", "prompt": "do it"})
        self.assertIn("not configured", str(ctx.exception))

    def test_configured_posts_only_id_prompt_repo_no_secrets(self):
        seen = {}

        def fake_client(url, payload, headers):
            seen["url"] = url
            seen["payload"] = payload
            seen["headers"] = headers
            return {"status": "completed", "stdout": "done"}

        ex = CloudExecutor(
            endpoint="https://runner.example.invalid/run",
            http_client=fake_client,
            auth_header="X-Runner-Token",
            env={"CLOUD_RUNNER_AUTH_TOKEN": "test-token-value"})
        result = ex.execute({"id": "t9", "prompt": "do it",
                             "repo": "https://example.com/r.git"})
        self.assertEqual(result["status"], "completed")
        self.assertEqual(seen["url"], "https://runner.example.invalid/run")
        self.assertEqual(set(seen["payload"].keys()),
                         {"id", "prompt", "repo_url"})
        self.assertEqual(seen["payload"]["repo_url"],
                         "https://example.com/r.git")
        self.assertNotIn("test-token-value",
                         json.dumps(seen["payload"]))
        self.assertEqual(seen["headers"]["X-Runner-Token"],
                         "test-token-value")
        # dispatch_task 是 worker 走的入口，必须与 execute 同结果。
        result2 = ex.dispatch_task({"id": "t9", "prompt": "do it"})
        self.assertEqual(result2["status"], "completed")

    def test_failed_status_maps_to_failed(self):
        ex = CloudExecutor(
            endpoint="https://runner.example.invalid/run",
            http_client=lambda u, p, h: {"status": "failed",
                                          "stdout": "boom"})
        result = ex.execute({"id": "t1", "prompt": "x"})
        self.assertEqual(result["status"], "failed")


class AutoPrTestCase(_ServerCase):
    def _pr(self, task_id, key):
        return self._req("POST", "/tasks/%s/pr" % task_id, {},
                         key)

    def test_auto_pr_disabled_skips_and_stays_merged(self):
        status, proj = self._req(
            "POST", "/projects", {"name": "p-off"}, self.owner_key)
        self.assertEqual(status, 201, proj)
        self._insert_task("toff", "merged", project_id=proj["id"])
        status, body = self._pr("toff", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertTrue(body["pr_skipped"])
        self.assertEqual(body["reason"], "auto_pr_disabled")
        self.assertEqual(self._task_row("toff")[0], "merged")

    def test_auto_pr_illegal_value_rejected(self):
        status, body = self._req(
            "POST", "/projects", {"name": "p-bad", "auto_pr": 2},
            self.owner_key)
        self.assertEqual(status, 400, body)

    def test_non_merged_returns_409(self):
        self._insert_task("tq", "queued")
        status, body = self._pr("tq", self.owner_key)
        self.assertEqual(status, 409, body)

    def test_viewer_cannot_pr(self):
        self._insert_task("tm", "merged")
        status, body = self._pr("tm", self.viewer_key)
        self.assertEqual(status, 403, body)

    def test_no_worktree_skips(self):
        self._insert_task("tnw", "merged")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tnw", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "no_worktree")
        self.assertEqual(self._task_row("tnw")[0], "merged")

    def test_no_remote_skips(self):
        self._insert_task("tnr", "merged")
        self._worktree("tnr")
        self.server.pr_git_runner = self._git_runner(remote_fail=True)
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tnr", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "no_remote")

    def test_no_gh_skips(self):
        self._insert_task("tngh", "merged")
        self._worktree("tngh")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = ""
        status, body = self._pr("tngh", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "no_gh")

    def test_gh_not_authenticated_skips(self):
        self._insert_task("tna", "merged")
        self._worktree("tna")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner(authed=False)
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tna", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["reason"], "gh_not_authenticated")

    def test_push_failure_returns_502_and_stays_merged(self):
        self._insert_task("tpf", "merged")
        self._worktree("tpf")
        self.server.pr_git_runner = self._git_runner(push_fail=True)
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tpf", self.owner_key)
        self.assertEqual(status, 502, body)
        self.assertEqual(body["step"], "push")
        self.assertEqual(self._task_row("tpf")[0], "merged")

    def test_success_writes_back_pr_and_audit(self):
        self._insert_task("tok", "merged")
        self._worktree("tok")
        self.server.pr_git_runner = self._git_runner()
        self.server.pr_gh_runner = self._gh_runner()
        self.server.pr_gh_which = "gh"
        status, body = self._pr("tok", self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["pr_url"],
                         "https://example.com/repo/pull/42")
        self.assertEqual(body["pr_number"], 42)
        self.assertEqual(body["branch"], "host/tok")
        row = self._task_row("tok")
        self.assertEqual(row[0], "merged")
        self.assertEqual(row[1], "https://example.com/repo/pull/42")
        self.assertEqual(row[2], 42)
        self.assertIn("pr_created", self._audit_actions())

    def test_pr_module_auto_pr_gate(self):
        with self.assertRaises(PrSkipped) as ctx:
            create_task_pr("x", 0)
        self.assertEqual(ctx.exception.reason, "auto_pr_disabled")


class ProjectTaskFieldTestCase(_ServerCase):
    def test_create_task_executor_validation(self):
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "x",
             "executor": "bogus"},
            self.owner_key)
        self.assertEqual(status, 400, body)
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "x",
             "executor": "docker"},
            self.owner_key)
        self.assertEqual(status, 201, body)
        self.assertEqual(body["executor"], "docker")


class WorkerTablePrecisionTestCase(unittest.TestCase):
    def test_workers_time_columns_are_double_precision(self):
        # 回归：2026-10-26 主机实测，workers.last_seen 用 REAL 时，
        # Postgres 的 float4 把 epoch 秒舍入到 128 秒一档，存活判断
        # 误差可达 ±64 秒、甚至出现负 age。新表必须是 DOUBLE PRECISION。
        import sqlite3 as _sqlite3
        tmp = tempfile.TemporaryDirectory()
        db_path = os.path.join(tmp.name, "w.db")
        from cloud.workers import ensure_workers_table
        ensure_workers_table(db_path)
        conn = _sqlite3.connect(db_path)
        try:
            rows = conn.execute("PRAGMA table_info(workers)").fetchall()
        finally:
            conn.close()
            tmp.cleanup()
        by_name = {r[1]: r[2] for r in rows}
        self.assertIn("DOUBLE", by_name["last_seen"])
        self.assertIn("DOUBLE", by_name["created_at"])


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