Worker concurrency (capacity enforced), survival (heartbeat/sweep/drain), retry endpoint

1. Concurrency: capacity is now a real limit, not display-only.
   run_all with capacity>1 runs a manual thread pool (threading,
   not concurrent.futures: a test in this repo puts cloud/ on
   sys.path, shadowing stdlib queue and breaking
   concurrent.futures' SimpleQueue import). Claim stays CAS, so
   two threads racing one task still yield one owner; max_tasks
   counts dispatched+running to avoid overshoot; stop signals
   stop new claims only, running tasks are drained by run_loop.
   Running state moved from single fields to a _running dict
   (proc/container per task) so cancel kills the right task.
   Docker timeout is now a per-call argument, never a mutated
   shared executor field. CostStore opens a fresh connection per
   record/report (shared sqlite connection is thread-bound and
   raised ProgrammingError under the pool; :memory: keeps the
   shared connection because a fresh one is an empty DB).
   Final status write is CAS on status='running' so a background
   thread cannot overwrite an interrupted task with completed.
2. Survival: run_loop starts a daemon heartbeat thread (every
   30s, refreshes running tasks' last_heartbeat) and sweeps stale
   tasks every round (threshold 300s, best-effort); live tasks
   never go stale now, dead workers' tasks are reclaimed without
   manual /admin/sweep. On exit, run_loop waits up to
   drain_timeout (CLI --drain-timeout, default 60s, inside
   systemd's 90s stop window) for running tasks before killing
   leftovers and marking interrupted.
3. Retry: POST /tasks/{id}/retry (maintainer+, 404/403/409):
   failed/interrupted/merge-failed go back to queued with
   result/error/model cleared, prompt/model/deps kept, CAS
   update, audit task_retried.
Tests: 5 concurrency tests (barrier overlap, serial capacity=1,
max_tasks, wait_for_running, CAS final write), 3 retry tests;
full suites pass.

cloud/tests/test_docker_sandbox.py

@@ -289,8 +289,7 @@ class WorkerDockerTestCase(unittest.TestCase):
             worker = Worker(self.db_path, "/fake/host.py",
                             env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex)
             # 模拟 dispatch 阻塞期间：worker 已把容器名记下来。
-            worker._running_task_id = "d3"
-            worker._running_container = "codex-d3"
+            worker._running_set("d3", container="codex-d3")
             conn = sqlite3.connect(self.db_path)
             try:
                 conn.execute(

cloud/tests/test_server_retry.py

@@ -0,0 +1,136 @@
+"""失败任务重跑端点测试：POST /tasks/{id}/retry。
+
+只测状态机本身，不执行真实任务：失败/中断任务重回 queued，
+上一轮结果清空，prompt 与模型原样保留；权限与状态码语义
+与 merge 端点对齐（maintainer 以上、404/401/403/409）。
+"""
+
+import http.client
+import json
+import os
+import sqlite3
+import sys
+import tempfile
+import threading
+import unittest
+from pathlib import Path
+
+try:
+    from cloud.auth import AuthStore
+    from cloud.server import create_server
+except ImportError:
+    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+    from cloud.auth import AuthStore
+    from cloud.server import create_server
+
+
+class RetryTestCase(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.contrib_id, self.contrib_key = self.auth.create_user("contrib")
+        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": "demo"}, self.owner_key)
+        self.assertEqual(status, 201)
+        self.auth.set_member(
+            self.project["id"], self.owner_id, self.contrib_id, "contributor")
+
+    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 _failed_task(self):
+        status, task = self._req(
+            "POST", "/tasks",
+            {"project_id": self.project["id"], "prompt": "x"},
+            self.owner_key)
+        self.assertEqual(status, 201)
+        conn = sqlite3.connect(self.db_path)
+        try:
+            conn.execute(
+                "UPDATE tasks SET status = 'failed',"
+                " error_text = 'old boom', actual_model = 'm1'"
+                " WHERE id = ?",
+                (task["id"],))
+            conn.commit()
+        finally:
+            conn.close()
+        return task
+
+    def test_retry_failed_by_owner_returns_queued_and_clears_result(self):
+        task = self._failed_task()
+        status, body = self._req(
+            "POST", "/tasks/%s/retry" % task["id"], {}, self.owner_key)
+        self.assertEqual(status, 200)
+        self.assertEqual(body["status"], "queued")
+        self.assertIsNone(body.get("error_text"))
+        self.assertIsNone(body.get("actual_model"))
+        conn = sqlite3.connect(self.db_path)
+        try:
+            actions = [r[0] for r in conn.execute(
+                "SELECT action FROM audit_log WHERE project_id = ?",
+                (self.project["id"],)).fetchall()]
+        finally:
+            conn.close()
+        self.assertIn("task_retried", actions)
+
+    def test_retry_non_failed_409(self):
+        status, task = self._req(
+            "POST", "/tasks",
+            {"project_id": self.project["id"], "prompt": "x"},
+            self.owner_key)
+        self.assertEqual(status, 201)
+        status, body = self._req(
+            "POST", "/tasks/%s/retry" % task["id"], {}, self.owner_key)
+        self.assertEqual(status, 409)
+
+    def test_retry_by_contributor_forbidden_and_missing_404(self):
+        task = self._failed_task()
+        status, _ = self._req(
+            "POST", "/tasks/%s/retry" % task["id"], {}, self.contrib_key)
+        self.assertEqual(status, 403)
+        status, _ = self._req(
+            "POST", "/tasks/nope/retry", {}, self.owner_key)
+        self.assertEqual(status, 404)
+        status, _ = self._req(
+            "POST", "/tasks/%s/retry" % task["id"], {})
+        self.assertEqual(status, 401)
+
+
+if __name__ == "__main__":
+    unittest.main()

cloud/tests/test_worker_concurrency.py

@@ -0,0 +1,188 @@
+"""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()