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

"""新业务链路单测：计划模式、后续消息（排队/直接发送）、会话分支。

不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
SQLite；worker 侧注入假 runner（假 CLI）与假执行器；链路目录
经 server.trace_dir 注入临时目录。
"""

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.sandbox import DockerExecutor
    from cloud.server import create_server, 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.auth import AuthStore
    from cloud.sandbox import DockerExecutor
    from cloud.server import create_server, init_db
    from cloud.worker import Worker


def _completed_runner(cmd):
    """假 CLI：dispatch/run 一律成功，打印 host.py 风格的 JSON。"""
    if "dispatch" in cmd:
        name = cmd[cmd.index("--name") + 1]
    elif "run" in cmd:
        name = cmd[cmd.index("--task") + 1]
    else:
        name = "x"
    return SimpleNamespace(
        returncode=0,
        stdout=json.dumps({"task": name, "status": "completed",
                           "text": "done"}),
        stderr="")


class ServerLinksTestCase(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.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()
        self.trace_dir = os.path.join(self.tmp.name, "host-traces")
        os.makedirs(self.trace_dir, exist_ok=True)
        self.server.trace_dir = self.trace_dir
        status, self.project = self._req(
            "POST", "/projects", {"name": "demo"}, self.owner_key)
        self.assertEqual(status, 201)

    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):
        conn = http.client.HTTPConnection(
            "127.0.0.1", self.port, timeout=10)
        hdrs = {}
        payload = None
        if body is not None:
            payload = json.dumps(body).encode("utf-8")
            hdrs["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 _create_task(self, **extra):
        body = {"project_id": self.project["id"], "prompt": "做一个加法函数"}
        body.update(extra)
        status, task = self._req("POST", "/tasks", body, self.owner_key)
        self.assertEqual(status, 201, task)
        return task

    def _set_status(self, task_id, status):
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
                         (status, task_id))
            conn.commit()
        finally:
            conn.close()

    # -- 计划模式 -------------------------------------------------------
    def test_plan_full_conflict_400(self):
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "p",
             "plan": 1, "approval": "full"},
            self.owner_key)
        self.assertEqual(status, 400)

    def test_plan_goal_conflict_400(self):
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "p",
             "plan": 1, "verify_command": "true"},
            self.owner_key)
        self.assertEqual(status, 400)

    def test_plan_create_roundtrip(self):
        task = self._create_task(plan=1, approval="ask")
        self.assertEqual(task.get("plan"), 1)
        status, got = self._req(
            "GET", "/tasks/" + task["id"], key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(got.get("plan"), 1)

    # -- 后续消息 -------------------------------------------------------
    def test_followup_validation(self):
        status, _ = self._req(
            "POST", "/tasks/nope/followups",
            {"body": "hi"}, self.owner_key)
        self.assertEqual(status, 404)
        task = self._create_task()
        status, _ = self._req(
            "POST", "/tasks/%s/followups" % task["id"],
            {"body": "   "}, self.owner_key)
        self.assertEqual(status, 400)
        status, _ = self._req(
            "POST", "/tasks/%s/followups" % task["id"],
            {"body": "hi", "mode": "bogus"}, self.owner_key)
        self.assertEqual(status, 400)
        # 终态任务拒绝后续消息。
        self._set_status(task["id"], "completed")
        status, _ = self._req(
            "POST", "/tasks/%s/followups" % task["id"],
            {"body": "hi"}, self.owner_key)
        self.assertEqual(status, 409)
        # goal 任务拒绝后续消息。
        goal = self._create_task(verify_command="true")
        status, _ = self._req(
            "POST", "/tasks/%s/followups" % goal["id"],
            {"body": "hi"}, self.owner_key)
        self.assertEqual(status, 409)

    def test_followup_create_and_list(self):
        task = self._create_task()
        status, fu1 = self._req(
            "POST", "/tasks/%s/followups" % task["id"],
            {"body": "先排队"}, self.owner_key)
        self.assertEqual(status, 201, fu1)
        self.assertEqual(fu1["status"], "pending")
        status, fu2 = self._req(
            "POST", "/tasks/%s/followups" % task["id"],
            {"body": "直接发", "mode": "direct"}, self.owner_key)
        self.assertEqual(status, 201, fu2)
        status, listing = self._req(
            "GET", "/tasks/%s/followups" % task["id"], key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual([f["body"] for f in listing["followups"]],
                         ["先排队", "直接发"])

    def test_followup_direct_appends_trace(self):
        task = self._create_task()
        status, _ = self._req(
            "POST", "/tasks/%s/followups" % task["id"],
            {"body": "立刻看到我", "mode": "direct"}, self.owner_key)
        self.assertEqual(status, 201)
        path = os.path.join(self.trace_dir, task["id"] + ".jsonl")
        with open(path, encoding="utf-8") as f:
            entries = [json.loads(line) for line in f if line.strip()]
        self.assertEqual(len(entries), 1)
        self.assertEqual(entries[0]["phase"], "followup-direct")
        self.assertIn("立刻看到我", entries[0]["text"])

    # -- 会话分支 -------------------------------------------------------
    def test_fork_terminal_inherits_fields(self):
        src = self._create_task(model="gpt-5.6-sol", effort="high",
                                approval="ask", plan=1)
        self._set_status(src["id"], "completed")
        status, forked = self._req(
            "POST", "/tasks/%s/fork" % src["id"],
            {"message": "换个方向继续"}, self.owner_key)
        self.assertEqual(status, 201, forked)
        self.assertEqual(forked["forked_from"], src["id"])
        self.assertEqual(forked["status"], "queued")
        self.assertEqual(forked["prompt"], "换个方向继续")
        self.assertEqual(forked["model"], "gpt-5.6-sol")
        self.assertEqual(forked["effort"], "high")
        self.assertEqual(forked["approval"], "ask")
        self.assertEqual(forked["plan"], 1)
        self.assertIsNone(forked.get("verify_command"))

    def test_fork_validation(self):
        task = self._create_task()
        # 运行中（queued）不能分支。
        status, _ = self._req(
            "POST", "/tasks/%s/fork" % task["id"],
            {"message": "x"}, self.owner_key)
        self.assertEqual(status, 409)
        self._set_status(task["id"], "completed")
        status, _ = self._req(
            "POST", "/tasks/%s/fork" % task["id"],
            {"message": "   "}, self.owner_key)
        self.assertEqual(status, 400)
        status, _ = self._req(
            "POST", "/tasks/nope/fork",
            {"message": "x"}, self.owner_key)
        self.assertEqual(status, 404)


def _insert_task(db_path, task_id, prompt="do work", status="queued",
                 project_id="p1", **extra):
    now = time.time()
    cols = ["id", "project_id", "prompt", "status", "idempotency_key",
            "depends_on", "priority", "created_by", "created_at",
            "updated_at", "last_heartbeat"]
    vals = [task_id, project_id, prompt, status, None, "[]", 0, "tester",
            now, now, None]
    for k, v in extra.items():
        cols.append(k)
        vals.append(v)
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT INTO tasks(%s) VALUES(%s)"
            % (", ".join(cols), ", ".join("?" * len(cols))), vals)
        conn.commit()
    finally:
        conn.close()


def _insert_followup(db_path, task_id, body, mode, created_at,
                     project_id="p1", status="pending"):
    # id 是 TEXT 主键（服务端生成 uuid）：不传 id 会插成 NULL，
    # CAS 认领 WHERE id = NULL 永远匹配不到，排空会静默跳过。
    import uuid
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT INTO task_followups(id, task_id, project_id, body,"
            " mode, status, created_by, created_at, updated_at)"
            " VALUES(?, ?, ?, ?, ?, ?, 'tester', ?, ?)",
            (uuid.uuid4().hex, task_id, project_id, body, mode, status,
             created_at, created_at))
        conn.commit()
    finally:
        conn.close()


def _row(db_path, table, task_id=None, fu_id=None):
    conn = sqlite3.connect(db_path)
    try:
        if fu_id is not None:
            row = conn.execute(
                "SELECT * FROM %s WHERE id = ?" % table,
                (fu_id,)).fetchone()
        else:
            row = conn.execute(
                "SELECT * FROM %s WHERE id = ?" % table,
                (task_id,)).fetchone()
        if row is None:
            return None
        cols = [d[0] for d in conn.execute(
            "SELECT * FROM %s LIMIT 0" % table).description]
        return dict(zip(cols, row))
    finally:
        conn.close()


def _followup_statuses(db_path, task_id):
    conn = sqlite3.connect(db_path)
    try:
        rows = conn.execute(
            "SELECT body, status FROM task_followups"
            " WHERE task_id = ? ORDER BY id", (task_id,)).fetchall()
    finally:
        conn.close()
    return dict(rows)


class WorkerLinksTestCase(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"}
        self.calls = []

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

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

    def _recording(self, inner):
        def fake(cmd):
            self.calls.append(list(cmd))
            return inner(cmd)
        return fake

    # -- 任务级模型优先 -------------------------------------------------
    def test_task_model_beats_worker_model(self):
        _insert_task(self.db_path, "T", model="gpt-5.6-sol")
        worker = self._worker(runner=self._recording(_completed_runner),
                              model="gpt-5.6-terra")
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        dispatch = [c for c in self.calls if "dispatch" in c][0]
        self.assertEqual(dispatch[dispatch.index("--model") + 1],
                         "gpt-5.6-sol")
        self.assertEqual(_row(self.db_path, "tasks", "T")["actual_model"],
                         "gpt-5.6-sol")

    # -- 计划模式命令 ---------------------------------------------------
    def test_plan_local_command_readonly(self):
        _insert_task(self.db_path, "T", plan=1)
        worker = self._worker(runner=self._recording(_completed_runner))
        worker.run_one()
        dispatch = [c for c in self.calls if "dispatch" in c][0]
        self.assertIn("--sandbox", dispatch)
        self.assertEqual(dispatch[dispatch.index("--sandbox") + 1],
                         "read-only")

    def test_plan_local_command_absent_without_plan(self):
        _insert_task(self.db_path, "T")
        worker = self._worker(runner=self._recording(_completed_runner))
        worker.run_one()
        dispatch = [c for c in self.calls if "dispatch" in c][0]
        self.assertNotIn("--sandbox", dispatch)

    def test_plan_docker_inner_command_readonly(self):
        ex = DockerExecutor(config_path="/fake/cfg.json")
        cmd = ex._inner_host_command(
            "dispatch", "--name", "x", "--task", "hi", sandbox="read-only")
        self.assertIn("--sandbox", cmd)
        self.assertEqual(cmd[cmd.index("--sandbox") + 1], "read-only")
        cmd2 = ex._inner_host_command(
            "dispatch", "--name", "x", "--task", "hi")
        self.assertNotIn("--sandbox", cmd2)

    # -- 后续消息排空 ---------------------------------------------------
    def test_followup_drain_order_direct_first(self):
        _insert_task(self.db_path, "T", prompt="orig")
        now = time.time()
        _insert_followup(self.db_path, "T", "q1", "queued", now)
        _insert_followup(self.db_path, "T", "d1", "direct", now + 1)
        _insert_followup(self.db_path, "T", "q2", "queued", now + 2)

        def fake(cmd):
            if "dispatch" in cmd:
                return SimpleNamespace(
                    returncode=1, stdout="",
                    stderr="Traceback: task 已存在：T，先换名或 discard 它")
            return _completed_runner(cmd)

        worker = self._worker(runner=self._recording(fake))
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        prompts = []
        for c in self.calls:
            if "run" in c:
                prompts.append(c[c.index("--prompt") + 1])
        self.assertEqual(prompts, ["orig", "d1", "q1", "q2"])
        self.assertEqual(_followup_statuses(self.db_path, "T"),
                         {"q1": "done", "d1": "done", "q2": "done"})

    def test_followup_failure_stops_drain(self):
        _insert_task(self.db_path, "T", prompt="orig")
        now = time.time()
        _insert_followup(self.db_path, "T", "bad", "queued", now)
        _insert_followup(self.db_path, "T", "later", "queued", now + 1)

        def fake(cmd):
            if "dispatch" in cmd and "bad" in cmd:
                return SimpleNamespace(returncode=1, stdout="",
                                       stderr="boom")
            return _completed_runner(cmd)

        worker = self._worker(runner=self._recording(fake))
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertEqual(_followup_statuses(self.db_path, "T"),
                         {"bad": "failed", "later": "pending"})

    def test_goal_task_does_not_drain(self):
        _insert_task(self.db_path, "T", prompt="orig",
                     verify_command="true")

        def fake(cmd):
            if "goal" in cmd:
                # goal 结论行：host.py goal 两种结局都 exit 0，
                # 结论靠这行解析，没有它会被判无结论。
                return SimpleNamespace(returncode=0,
                                       stdout="goal met after rounds\n",
                                       stderr="")
            return _completed_runner(cmd)

        _insert_followup(self.db_path, "T", "later", "queued", time.time())
        worker = self._worker(runner=self._recording(fake))
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertEqual(_followup_statuses(self.db_path, "T"),
                         {"later": "pending"})

    # -- 会话分支 -------------------------------------------------------
    def test_fork_runs_first_round_once(self):
        _insert_task(self.db_path, "T", prompt="branch msg",
                     forked_from="SRC")
        forks = []

        def fake(cmd):
            if "fork" in cmd:
                forks.append(list(cmd))
                return SimpleNamespace(returncode=0,
                                       stdout="forked SRC -> T\n",
                                       stderr="")
            if "dispatch" in cmd:
                return SimpleNamespace(
                    returncode=1, stdout="",
                    stderr="Traceback: task 已存在：T，先换名或 discard 它")
            return _completed_runner(cmd)

        worker = self._worker(runner=self._recording(fake))
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertEqual(len(forks), 1)
        fork = forks[0]
        self.assertIn("SRC", fork)
        self.assertIn("T", fork)
        runs = [c for c in self.calls if "run" in c]
        self.assertEqual(len(runs), 1)
        self.assertEqual(runs[0][runs[0].index("--prompt") + 1],
                         "branch msg")

    def test_fork_failure_not_retried(self):
        _insert_task(self.db_path, "T", prompt="branch msg",
                     forked_from="SRC")
        forks = []

        def fake(cmd):
            if "fork" in cmd:
                forks.append(list(cmd))
                return SimpleNamespace(returncode=0,
                                       stdout="forked SRC -> T\n",
                                       stderr="")
            if "dispatch" in cmd:
                return SimpleNamespace(
                    returncode=1, stdout="",
                    stderr="Traceback: task 已存在：T，先换名或 discard 它")
            return SimpleNamespace(returncode=1, stdout="",
                                   stderr="model boom")

        worker = self._worker(runner=self._recording(fake))
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        # range(0) 跳过重试循环：fork 只调一次，绝不重复分叉。
        self.assertEqual(len(forks), 1)


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