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/cost.py

@@ -76,6 +76,41 @@ class CostStore:
         )
         self._conn.commit()
 
+    def _new_conn(self):
+        """每次写/读开一个新连接，用完即关。
+
+        旧实现在 __init__ 里建一条连接并终身共用，单线程没问题；
+        worker 并发后多个线程共用同一条 sqlite 连接会直接报
+        ProgrammingError（连接是线程绑定的）。其它 store（Worker
+        的 _set_status、_audit 等）本来就是每次开新连接，这里
+        向它们看齐。_conn 仍保留（建表 + 旧测试直接用它读），
+        但 record/_report 不再走它。
+        例外：db_path 为 :memory: 时新连接是另一个空库，必须
+        复用 _conn（:memory: 只用于单线程单测，并发 worker
+        从不用它）。
+        """
+        if self.db_path == ":memory:":
+            return self._conn
+        conn = _db_connect(self.db_path)
+        conn.execute(
+            """CREATE TABLE IF NOT EXISTS cost_records (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                task_id TEXT,
+                project_id TEXT,
+                model TEXT,
+                input_tokens INTEGER,
+                output_tokens INTEGER,
+                total_tokens INTEGER,
+                usd_micros INTEGER,
+                source TEXT,
+                fallback_count INTEGER DEFAULT 0,
+                day TEXT,
+                created_at REAL
+            )"""
+        )
+        conn.commit()
+        return conn
+
     def record(self, task_id, project_id, model, input_tokens=0,
                output_tokens=0, usd_micros=None, source="actual",
                fallback_count=0):
@@ -93,18 +128,24 @@ class CostStore:
         total_tokens = input_tokens + output_tokens
         created_at = time.time()
         day = datetime.fromtimestamp(created_at, tz=timezone.utc).strftime("%Y-%m-%d")
-        cursor = self._conn.execute(
-            """INSERT INTO cost_records
-               (task_id, project_id, model, input_tokens, output_tokens,
-                total_tokens, usd_micros, source, fallback_count, day,
-                created_at)
-               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
-            (task_id, project_id, model, input_tokens, output_tokens,
-             total_tokens, usd_micros, source, int(fallback_count or 0),
-             day, created_at),
-        )
-        self._conn.commit()
-        return cursor.lastrowid
+        conn = self._new_conn()
+        try:
+            cursor = conn.execute(
+                """INSERT INTO cost_records
+                   (task_id, project_id, model, input_tokens, output_tokens,
+                    total_tokens, usd_micros, source, fallback_count, day,
+                    created_at)
+                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                (task_id, project_id, model, input_tokens, output_tokens,
+                 total_tokens, usd_micros, source, int(fallback_count or 0),
+                 day, created_at),
+            )
+            conn.commit()
+            return cursor.lastrowid
+        finally:
+            # :memory: 时 conn 就是共享的 _conn，绝不能关掉它。
+            if conn is not self._conn:
+                conn.close()
 
     def record_estimated(self, task_id, project_id, model, prompt_text,
                          output_text):
@@ -123,19 +164,24 @@ class CostStore:
         )
 
     def _report(self, where, params):
-        row = self._conn.execute(
-            """SELECT
-                 COALESCE(SUM(total_tokens), 0),
-                 COALESCE(SUM(usd_micros), 0),
-                 COUNT(*),
-                 COALESCE(SUM(fallback_count), 0),
-                 COALESCE(SUM(CASE WHEN source = 'actual'
-                                   THEN total_tokens ELSE 0 END), 0),
-                 COALESCE(SUM(CASE WHEN source = 'estimated'
-                                   THEN total_tokens ELSE 0 END), 0)
-               FROM cost_records WHERE %s""" % where,
-            params,
-        ).fetchone()
+        conn = self._new_conn()
+        try:
+            row = conn.execute(
+                """SELECT
+                     COALESCE(SUM(total_tokens), 0),
+                     COALESCE(SUM(usd_micros), 0),
+                     COUNT(*),
+                     COALESCE(SUM(fallback_count), 0),
+                     COALESCE(SUM(CASE WHEN source = 'actual'
+                                       THEN total_tokens ELSE 0 END), 0),
+                     COALESCE(SUM(CASE WHEN source = 'estimated'
+                                       THEN total_tokens ELSE 0 END), 0)
+                   FROM cost_records WHERE %s""" % where,
+                params,
+            ).fetchone()
+        finally:
+            if conn is not self._conn:
+                conn.close()
         total_tokens, usd_micros, turns, fallback_count, actual, estimated = row
         return {
             "total_tokens": int(total_tokens),

cloud/sandbox.py

@@ -420,10 +420,17 @@ class DockerExecutor(Executor):
         return cmd
 
     # -- runner 调用 / 清理 --------------------------------------------
-    def _call_runner(self, cmd):
-        """调 runner：优先传 timeout，假 runner 只接受 cmd 时回退。"""
+    def _call_runner(self, cmd, timeout=None):
+        """调 runner：优先传 timeout，假 runner 只接受 cmd 时回退。
+
+        timeout 参数只影响本次调用（并发下两个任务各用各的
+        超时）；不传时用实例的 self.timeout。绝不能在调用方
+        临时改 self.timeout 再恢复，那在并发下会互相覆盖。
+        """
         try:
-            return self.runner(cmd, timeout=self.timeout)
+            return self.runner(
+                cmd,
+                timeout=self.timeout if timeout is None else timeout)
         except TypeError:
             return self.runner(cmd)
 
@@ -493,20 +500,20 @@ class DockerExecutor(Executor):
                 continue
         return ok
 
-    def _run_in_container(self, task, inner_cmd):
+    def _run_in_container(self, task, inner_cmd, timeout=None):
         """在容器里跑一轮内层命令，返回 LocalExecutor 同款结果 dict。
 
         超时：runner 抛 subprocess.TimeoutExpired（或同名异常）
         时先 kill+rm 容器，再返回 failed（stdout 截断 + timeout
         说明，returncode 为 None）；其它 runner 异常也如实返回
-        failed，不编造成功。
+        failed，不编造成功。timeout 只影响本次调用，不改实例。
         """
         task_id = task.get("id") if isinstance(task, dict) else None
         if task_id is None:
             raise ValueError("任务缺少 id")
         cmd = self.build_run_command(task, inner_cmd)
         try:
-            result = self._call_runner(cmd)
+            result = self._call_runner(cmd, timeout=timeout)
         except Exception as exc:
             if exc.__class__.__name__ == "TimeoutExpired":
                 self.kill_container(task)
@@ -529,7 +536,7 @@ class DockerExecutor(Executor):
         parsed["container"] = self.container_name(task)
         return parsed
 
-    def dispatch_task(self, task) -> dict:
+    def dispatch_task(self, task, timeout=None) -> dict:
         """容器内调 host.py dispatch：建 host 任务并跑一轮 turn。"""
         if not isinstance(task, dict) or task.get("id") is None:
             raise ValueError("任务缺少 id")
@@ -541,9 +548,9 @@ class DockerExecutor(Executor):
         repo = task.get("repo")
         if repo:
             inner += ["--repo", repo]
-        return self._run_in_container(task, inner)
+        return self._run_in_container(task, inner, timeout=timeout)
 
-    def run_task(self, task) -> dict:
+    def run_task(self, task, timeout=None) -> dict:
         """容器内调 host.py run：在已有 host 任务上再跑一轮 turn。"""
         if not isinstance(task, dict) or task.get("id") is None:
             raise ValueError("任务缺少 id")
@@ -552,7 +559,7 @@ class DockerExecutor(Executor):
         inner = self._inner_host_command(
             "run", "--task", self._host_name(task),
             "--prompt", task["prompt"], model=task.get("model"))
-        return self._run_in_container(task, inner)
+        return self._run_in_container(task, inner, timeout=timeout)
 
     def execute(self, task) -> dict:
         """只读执行：容器内调 host.py status，不直接跑模型回合。

cloud/server.py

@@ -624,6 +624,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
                 self._merge_task(user_id, rest[: -len("/merge")])
             elif rest.endswith("/pr") and method == "POST":
                 self._task_pr(user_id, rest[: -len("/pr")])
+            elif rest.endswith("/retry") and method == "POST":
+                self._retry_task(user_id, rest[: -len("/retry")])
             elif rest.endswith("/heartbeat") and method == "POST":
                 self._heartbeat_task(user_id, rest[: -len("/heartbeat")])
             else:
@@ -1321,6 +1323,67 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
             store.audit(user_id, task["project_id"], "merge_task", task_id)
         self._send_json(200, task)
 
+    def _retry_task(self, user_id, task_id):
+        """把终态失败任务重新排队：只改状态，不直接执行。
+
+        规则：任务必须存在（404）、调用方必须是项目 maintainer
+        以上（approve_merge，403——重跑要再花模型钱，不能让
+        只有读权限或普通贡献者随便点）、任务必须是 failed /
+        interrupted / merge-failed（其它状态一律 409：running
+        本来就在跑，queued 本来就在等，completed/merged 已有
+        结果，重跑应是新任务而不是覆盖结果）。
+        成功后状态回 queued，上一轮的 result_text/error_text/
+        actual_model 清空（否则新一轮的结果会和旧结果混在一起
+        无法分辨）；prompt、模型、优先级、依赖原样保留，依赖
+        已完成时重跑任务会立刻就绪。写审计 task_retried。
+        更新用 CAS（WHERE status IN (...)），避免与正在完成的
+        回合互相覆盖：若 CAS 没改到行，说明状态已变，返回
+        409 而不是假装成功。
+        """
+        row = self._fetch_task(task_id)
+        if row is None:
+            self._send_json(404, {"error": "task not found"})
+            return
+        task = _task_to_dict(row)
+        with _auth_store(self.server.db_path) as store:
+            if not store.check_permission(
+                    user_id, task["project_id"], "approve_merge"):
+                self._send_json(403, {"error": "forbidden"})
+                return
+        if task["status"] not in ("failed", "interrupted", "merge-failed"):
+            self._send_json(409, {
+                "error": "only failed/interrupted tasks can be retried",
+                "status": task["status"],
+            })
+            return
+        now = time.time()
+        conn = self._db()
+        try:
+            cur = conn.execute(
+                "UPDATE tasks SET status = 'queued', result_text = NULL,"
+                " error_text = NULL, actual_model = NULL, updated_at = ?"
+                " WHERE id = ? AND status IN"
+                " ('failed', 'interrupted', 'merge-failed')",
+                (now, task_id),
+            )
+            conn.commit()
+            if cur.rowcount != 1:
+                self._send_json(409, {
+                    "error": "task status changed, retry aborted",
+                    "status": task["status"],
+                })
+                return
+            task = _task_to_dict(
+                conn.execute(
+                    "SELECT * FROM tasks WHERE id = ?", (task_id,)
+                ).fetchone()
+            )
+        finally:
+            conn.close()
+        with _auth_store(self.server.db_path) as store:
+            store.audit(user_id, task["project_id"], "task_retried", task_id)
+        self._send_json(200, task)
+
     # -- worker 池 -----------------------------------------------------
     def _workers_list(self, user_id):
         """列出注册 worker：任何已认证用户可读（worker 是基础设施

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()

cloud/worker.py

@@ -117,6 +117,7 @@ import signal
 import socket
 import subprocess
 import sys
+import threading
 import time
 from types import SimpleNamespace
 
@@ -270,13 +271,21 @@ class Worker:
                 repo=repo)
         self._cost_store = cost_store
         self._cancel_requested = set()
-        # 真取消用：正在跑的子进程（Popen 对象）与任务 id。
-        # process_factory 可注入：单测用它返回一个假长跑子进程，
-        # 真实执行时走 LocalExecutor.popen_task（进程组启动）。
+        # 真取消用：process_factory 可注入（单测用它返回一个假长跑
+        # 子进程，真实执行时走 LocalExecutor.popen_task 进程组启动）。
         self.process_factory = process_factory
-        self._running_proc = None
-        self._running_task_id = None
-        self._running_container = None
+        # 并发运行状态：task_id -> {"proc": Popen|None,
+        # "container": 容器名|None}，受 _running_lock 保护。
+        # 旧版只有单任务三个字段（_running_proc 等），capacity=1
+        # 时够用；capacity>1 后多个任务同时跑，单字段会被后
+        # 一个任务覆盖，导致 cancel 杀错、退出时漏杀。所有读写
+        # 都走下面的 _running_set/_running_clear/_running_snapshot。
+        self._running = {}
+        self._running_lock = threading.Lock()
+        # DB 写锁：cost 记录共用一条连接（CostStore._conn），
+        # usage 计数也是共享变量；模型回合本身在子进程/容器里
+        # 跑不受这把锁影响，锁只覆盖毫秒级的写库操作。
+        self._db_lock = threading.RLock()
         self._loop_stop = False
         self.logger = get_logger("cloud.worker")
         # 任务级超时 / 失败重试（见模块 docstring 的生产级说明）。
@@ -469,6 +478,27 @@ class Worker:
                 except Exception:
                     pass
 
+    def _set_status_if_running(self, task_id, status):
+        """仅当任务仍为 running 时改状态，返回是否改成功。
+
+        并发 + 优雅排空下，后台线程可能在任务已被取消
+        （标 interrupted）后才跑完；此时最终写回绝不能把
+        interrupted 盖回 completed/failed。正常流程里任务
+        到最终写回时仍是 running，CAS 一定成功。
+        """
+        now = time.time()
+        conn = _db_connect(self.db_path)
+        try:
+            cur = conn.execute(
+                "UPDATE tasks SET status = ?, last_heartbeat = ?,"
+                " updated_at = ? WHERE id = ? AND status = 'running'",
+                (status, now, now, task_id),
+            )
+            conn.commit()
+            return cur.rowcount == 1
+        finally:
+            conn.close()
+
     def _write_execution_result(self, task_id, result_text=None,
                                 error_text=None, actual_model=None,
                                 started_at=None, finished_at=None,
@@ -597,6 +627,53 @@ class Worker:
                  if t.get("executor") in (None, "", mine)]
         return ready[0] if ready else None
 
+    # -- 并发运行状态 --------------------------------------------------
+    def _running_set(self, task_id, proc=None, container=None):
+        with self._running_lock:
+            entry = self._running.setdefault(task_id, {})
+            if proc is not None:
+                entry["proc"] = proc
+            if container is not None:
+                entry["container"] = container
+
+    def _running_clear(self, task_id):
+        with self._running_lock:
+            self._running.pop(task_id, None)
+
+    def _running_snapshot(self):
+        with self._running_lock:
+            return {tid: dict(entry)
+                    for tid, entry in self._running.items()}
+
+    def _wait_for_running(self, timeout=None):
+        """等 _running 变空；timeout 为 None 时无限等。
+
+        返回 True 表示已空，False 表示超时仍有任务在跑。
+        run_loop 退出时的优雅排空用它：先等任务自然结束，
+        超时后调用方再杀残留并标 interrupted。
+        """
+        deadline = None if timeout is None else time.time() + timeout
+        while True:
+            with self._running_lock:
+                if not self._running:
+                    return True
+            if deadline is not None and time.time() >= deadline:
+                return False
+            time.sleep(0.2)
+
+    def _heartbeat_running_tasks(self):
+        """给所有正在跑的任务刷 last_heartbeat。
+
+        run_loop 的心跳线程调它：任务在跑时心跳保持新鲜，
+        自动陈旧回收（sweep）才不会误杀活任务；worker 真死后
+        心跳停了，任务超过阈值才会被别的 worker 回收。
+        """
+        for task_id in self._running_snapshot():
+            try:
+                self.heartbeat(task_id)
+            except Exception:
+                pass
+
     # -- 取消（server cancel 语义，见模块 docstring）-------------------
     def cancel(self, task_id):
         """把 queued/running 任务标 interrupted；已终态或不存在返回 False。
@@ -608,26 +685,24 @@ class Worker:
         if task is None or task["status"] in TERMINAL_STATUSES:
             return False
         self._cancel_requested.add(task_id)
-        if self._running_task_id == task_id and self._running_proc is not None:
-            self._kill_proc_group(self._running_proc)
+        entry = self._running_snapshot().get(task_id) or {}
+        if entry.get("proc") is not None:
+            self._kill_proc_group(entry["proc"])
             try:
-                self._running_proc.wait(timeout=3)
+                entry["proc"].wait(timeout=3)
             except Exception:
                 pass
-            self._running_proc = None
-            self._running_task_id = None
         # Docker 取消：按任务派生的容器名真杀，不只改数据库状态。
-        # 同步路径（runner 阻塞在 docker run 里）下 _running_proc
-        # 为空，此时靠 _running_container 杀容器；kill 尽力而为，
-        # 失败也不影响后续的状态标记。
-        if self._running_task_id == task_id and \
-                self._running_container is not None and \
+        # 同步路径（runner 阻塞在 docker run 里）下 proc 为空，
+        # 此时靠容器名杀容器；kill 尽力而为，失败也不影响后续
+        # 的状态标记。
+        if entry.get("container") is not None and \
                 hasattr(self.executor, "kill_container"):
             try:
-                self.executor.kill_container(self._running_container)
+                self.executor.kill_container(entry["container"])
             except Exception:
                 pass
-            self._running_container = None
+        self._running_clear(task_id)
         conn = _db_connect(self.db_path)
         try:
             cur = conn.execute(
@@ -724,8 +799,7 @@ class Worker:
             proc = self.process_factory(task)
         else:
             proc = self.executor.popen_task(task, kind=kind)
-        self._running_proc = proc
-        self._running_task_id = task_id
+        self._running_set(task_id, proc=proc)
         try:
             self._set_child_pid(task_id, getattr(proc, "pid", None))
         except Exception:
@@ -754,8 +828,7 @@ class Worker:
                     pass
                 out, err = proc.communicate()
         finally:
-            self._running_proc = None
-            self._running_task_id = None
+            self._running_clear(task_id)
         if cancelled:
             return {"executor": "local", "task_id": task_id,
                     "status": "cancelled", "stdout": out or "",
@@ -839,28 +912,30 @@ class Worker:
         result = None
         error = None
         # Docker 同步路径：dispatch 阻塞在 docker run 里时，取消
-        # 线程靠 _running_container 按名杀容器（见 cancel()）。
+        # 线程靠容器名杀容器（见 cancel()）。容器名按任务派生，
+        # 与执行器本身无关，所以这里只记名不碰执行器状态。
         if hasattr(self.executor, "container_name"):
             try:
-                self._running_container = \
-                    self.executor.container_name(task)
-                self._running_task_id = task_id
+                self._running_set(
+                    task_id,
+                    container=self.executor.container_name(task))
             except Exception:
-                self._running_container = None
+                pass
         # Docker 超时取较小值：两个超时都不应静默延长另一个。
-        saved_docker_timeout = None
+        # 注意绝不能改共享的 executor.timeout（并发下两个任务会
+        # 互相覆盖对方的超时）；超时只作为本次调用的参数传下去。
+        effective_docker_timeout = None
         is_docker = isinstance(self.executor, DockerExecutor)
         if is_docker and self.task_timeout is not None:
-            saved_docker_timeout = self.executor.timeout
-            if saved_docker_timeout is None:
-                self.executor.timeout = self.task_timeout
+            if self.executor.timeout is None:
+                effective_docker_timeout = self.task_timeout
             else:
                 try:
-                    self.executor.timeout = min(
+                    effective_docker_timeout = min(
                         float(self.task_timeout),
-                        float(saved_docker_timeout))
+                        float(self.executor.timeout))
                 except (TypeError, ValueError):
-                    self.executor.timeout = self.task_timeout
+                    effective_docker_timeout = self.task_timeout
         try:
             if self._use_popen():
                 result = self._execute_via_process(
@@ -874,26 +949,36 @@ class Worker:
                         task, kind="run",
                         deadline_seconds=self.task_timeout)
             else:
-                result = self.executor.dispatch_task(task)
+                if is_docker:
+                    result = self.executor.dispatch_task(
+                        task, timeout=effective_docker_timeout)
+                else:
+                    result = self.executor.dispatch_task(task)
                 combined = (result.get("stdout", "") or "") + \
                     (result.get("stderr", "") or "")
                 if result["status"] != "completed" and \
                         _ALREADY_EXISTS_MARKER in combined:
-                    result = self.executor.run_task(task)
+                    if is_docker:
+                        result = self.executor.run_task(
+                            task, timeout=effective_docker_timeout)
+                    else:
+                        result = self.executor.run_task(task)
         except Exception as exc:  # runner 异常也如实写回 failed
             result = {"executor": "local", "task_id": task_id,
                       "status": "failed", "stdout": "", "returncode": None,
                       "host_status": None, "stderr": ""}
             error = str(exc)
         finally:
-            if is_docker:
-                self.executor.timeout = saved_docker_timeout
             # 同步路径执行结束：清掉取消用的容器标记（取消路径里
-            # 若已清过，这里再清一次无害）。
-            self._running_container = None
-            if self._running_task_id == task_id and \
-                    self._running_proc is None:
-                self._running_task_id = None
+            # 若已清过，这里再清一次无害）。进程模式的 proc 由
+            # _execute_via_process 自己的 finally 清，这里只清
+            # 容器项：_running_clear 整个删掉该任务项会连 proc
+            # 也一起丢，所以这里只在 proc 为空时清容器标记——
+            # 实际做法是：如果该任务项里没有 proc，就整项清除；
+            # 有 proc（进程模式）时不动它，等它自己的 finally。
+            entry = self._running_snapshot().get(task_id) or {}
+            if entry.get("proc") is None:
+                self._running_clear(task_id)
         # Docker 超时标记：执行器已在超时时 kill+rm 容器并返回
         # failed（stdout 含 timeout），这里把状态统一成 timeout，
         # 调用方据此置 failed 并记 worker_timeout 审计。
@@ -1185,29 +1270,38 @@ class Worker:
                         fallback_count = 1
                         self._audit(task["project_id"], "worker_fallback",
                                     "%s -> %s" % (self.model, nxt))
-        self._set_status(task_id, final)
+        final_written = self._set_status_if_running(task_id, final)
         finished_at = time.time()
         self._write_execution_result(
             task_id, result_text=stdout, error_text=error,
             actual_model=actual_model, started_at=started_at,
             finished_at=finished_at)
-        self.usage_turns += 1
-        # 成本：host 协议不暴露真实 token 数，用 cost.py 的估算结构
-        # 记 estimated，绝不标 actual；fallback_count 复用同一记录；
-        # 模型记最后真正跑的那个（fallback 成功时就是 fallback 模型）。
-        total = estimate_tokens(task.get("prompt") or "", stdout)
-        cost_id = self.costs.record(
-            task_id, task["project_id"], actual_model,
-            input_tokens=total, output_tokens=0, usd_micros=None,
-            source="estimated", fallback_count=fallback_count)
-        self._audit(task["project_id"],
-                    "worker_completed" if final == "completed"
-                    else "worker_failed", task_id)
-        log_event(self.logger,
-                  "worker_completed" if final == "completed"
-                  else "worker_failed",
-                  task_id=task_id, project_id=task["project_id"],
-                  status=final)
+        with self._db_lock:
+            self.usage_turns += 1
+            # 成本：host 协议不暴露真实 token 数，用 cost.py 的估算结构
+            # 记 estimated，绝不标 actual；fallback_count 复用同一记录；
+            # 模型记最后真正跑的那个（fallback 成功时就是 fallback 模型）。
+            # cost 记录共用一条连接，并发下必须持锁写，否则两个线程
+            # 共用同一 sqlite 连接会直接报错。
+            total = estimate_tokens(task.get("prompt") or "", stdout)
+            cost_id = self.costs.record(
+                task_id, task["project_id"], actual_model,
+                input_tokens=total, output_tokens=0, usd_micros=None,
+                source="estimated", fallback_count=fallback_count)
+        if final_written:
+            self._audit(task["project_id"],
+                        "worker_completed" if final == "completed"
+                        else "worker_failed", task_id)
+            log_event(self.logger,
+                      "worker_completed" if final == "completed"
+                      else "worker_failed",
+                      task_id=task_id, project_id=task["project_id"],
+                      status=final)
+        else:
+            # 已被取消/回收：DB 状态以 interrupted 为准，不再
+            # 补一条 worker_completed/failed 审计误导排查。
+            final = "interrupted"
+            self._audit(task["project_id"], "worker_interrupted", task_id)
         return {"dispatched": True, "task_id": task_id,
                 "host_name": self.executor._host_name(task)
                 if hasattr(self.executor, "_host_name") else task_id,
@@ -1221,22 +1315,122 @@ class Worker:
 
         每轮都重新读库，所以上一轮刚 completed 的任务会解锁它的
         依赖任务；每轮至少推进一个任务到终态，不会空转死循环。
+
+        capacity<=1 时走旧串行循环（行为与原来完全一致）。
+        capacity>1 时用线程池并行跑：同一 worker 内最多
+        capacity 个任务同时执行（capacity 从此是真实并发上限，
+        不再只是注册表展示）。认领本身是 CAS，两个线程同时
+        盯上同一任务时只有一个能拿到它，另一个的 run_one 返回
+        already_claimed，不计入结果、不死循环。依赖解锁靠
+        “等一个在跑任务完成后再认领”：pending 非空但暂无就绪
+        任务时，等一个完成再继续，而不是直接退出。
+        停机信号（_loop_stop/stop_event）只阻止新认领，不会
+        丢弃已提交的任务；排空等待由 run_loop 的 finally 负责。
         """
+        if self.capacity <= 1:
+            results = []
+            while max_tasks is None or len(results) < max_tasks:
+                result = self.run_one()
+                if not result["dispatched"]:
+                    break
+                results.append(result)
+            return results
+        # capacity>1：手动线程池（不用 concurrent.futures）。
+        # 原因：本仓库有个单测把 cloud/ 插进 sys.path 后
+        # `import queue` 会解析到 cloud/queue.py 而不是标准库，
+        # concurrent.futures 内部依赖标准库 queue.SimpleQueue，
+        # 在该单测之后导入它会直接 AttributeError；手动用
+        # threading.Condition + 共享计数实现线程池，避开这个坑。
+        # 池线程数 = capacity；每个池线程循环认领-执行，认领
+        # 本身是 CAS，两个线程盯上同一任务时只有一个能拿到。
+        # 依赖解锁：某线程暂无就绪任务时，若仍有线程在跑就等
+        # 条件变量（跑完的线程会 notify，可能解锁了依赖），
+        # 全部线程都空闲且无就绪任务时才退出。
+        # max_tasks 约束：已派发结果数 + 正在跑的任务数达到
+        # max_tasks 就不再认领，避免线程池提前提交导致超跑。
+        # 停机信号：只阻止新认领，不丢弃已在跑的任务；run_all
+        # 在停机时立刻返回（后台线程继续跑），run_loop 的
+        # finally 做带超时的排空，超时后才杀残留。
+        state_lock = threading.Condition()
+        shared = {"dispatched": 0, "active": 0, "stopped": False}
         results = []
-        while max_tasks is None or len(results) < max_tasks:
-            result = self.run_one()
-            if not result["dispatched"]:
+        results_lock = threading.Lock()
+
+        def _pool_thread():
+            while True:
+                with state_lock:
+                    if shared["stopped"] or self._loop_stop:
+                        return
+                    if max_tasks is not None and \
+                            shared["dispatched"] + shared["active"] >= max_tasks:
+                        return
+                    try:
+                        task = self.pick_next()
+                    except Exception as exc:
+                        # 认领读库失败（如缺表、锁冲突）不炸池线程：
+                        # 等一轮再试，与 run_loop 的容错一致。
+                        log_event(self.logger, "worker_pool_pick_error",
+                                  message="worker pool pick error: %s" % exc)
+                        state_lock.wait(timeout=0.2)
+                        continue
+                    if task is None:
+                        if shared["active"] > 0:
+                            state_lock.wait(timeout=0.2)
+                            continue
+                        return
+                    shared["active"] += 1
+                result = {"dispatched": False}
+                try:
+                    result = self.run_one(task)
+                except Exception as exc:
+                    # 单个任务的意外异常绝不能拖垮整个池（串行
+                    # 模式下 run_loop 也是这么容错的）；该任务已
+                    # 认领为 running，之后由自动陈旧回收处理。
+                    log_event(self.logger, "worker_pool_task_error",
+                              message="worker pool task error: %s" % exc)
+                finally:
+                    with state_lock:
+                        shared["active"] -= 1
+                        if result["dispatched"]:
+                            shared["dispatched"] += 1
+                        state_lock.notify_all()
+                if result["dispatched"]:
+                    with results_lock:
+                        results.append(result)
+
+        threads = [threading.Thread(target=_pool_thread, daemon=True)
+                   for _ in range(self.capacity)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            # 正常情况等全部池线程结束；停机信号来时不等了，
+            # 立刻返回给 run_loop 做排空（后台线程是 daemon，
+            # 且 run_loop finally 会杀残留，不会悬着进程）。
+            while t.is_alive():
+                t.join(timeout=0.2)
+                if self._loop_stop:
+                    break
+            if self._loop_stop:
                 break
-            results.append(result)
         return results
 
-    def run_loop(self, interval=300, max_seconds=None, stop_event=None):
+    def run_loop(self, interval=300, max_seconds=None, stop_event=None,
+                 drain_timeout=60, sweep_threshold=300):
         """常驻循环：每 interval 秒跑一轮 run_all，空闲也不退出。
 
         退出条件：stop_event 被置位、_loop_stop 被信号处理置位、
-        或 max_seconds 到期。退出时若有正在跑的任务（进程模式下
-        _running_task_id 非空），先杀进程组、再标 interrupted，
-        不让任务永远卡在 running。
+        或 max_seconds 到期。退出时先优雅排空：等正在跑的任务
+        自然结束（最多 drain_timeout 秒），超时后才杀残留进程组/
+        容器并标 interrupted——以前是退出时立刻杀，部署重启会
+        无谓中断一个 16 秒就能跑完的任务。
+
+        运行中保障（与并发一起上的两件套）：
+        - 心跳线程：每 30 秒给所有正在跑的任务刷 last_heartbeat，
+          否则长任务（>300 秒）在跑也会被陈旧回收误杀。
+        - 自动陈旧回收：每轮结束调一次 sweep_stale_tasks（阈值
+          sweep_threshold，默认 300 秒），把死 worker 遗留的
+          running 任务标 interrupted 并审计；活任务因心跳新鲜
+          绝不会被扫到。回收失败只记日志，不炸循环。
         """
         # 常驻启动清理兜底：清掉上一轮残留的 codex- 容器，
         # 失败只记日志（见 cleanup_stale_containers），不炸循环。
@@ -1249,6 +1443,18 @@ class Worker:
         self.register()
         started = time.time()
         all_results = []
+        # 心跳线程：daemon，run_loop 退出即停，不阻塞进程退出。
+        hb_stop = threading.Event()
+
+        def _hb_loop():
+            while not hb_stop.is_set() and not self._loop_stop:
+                if stop_event is not None and stop_event.is_set():
+                    break
+                self._heartbeat_running_tasks()
+                hb_stop.wait(30)
+
+        hb_thread = threading.Thread(target=_hb_loop, daemon=True)
+        hb_thread.start()
         try:
             while not self._loop_stop:
                 if stop_event is not None and stop_event.is_set():
@@ -1265,6 +1471,15 @@ class Worker:
                               message="worker_loop_error: %s" % exc)
                     results = []
                 all_results.extend(results)
+                # 自动陈旧回收：死 worker 的 running 任务在这里被
+                # 标 interrupted；活任务心跳新鲜，不会被扫到。
+                if sweep_threshold is not None:
+                    try:
+                        from cloud.server import sweep_stale_tasks
+                        sweep_stale_tasks(self.db_path, sweep_threshold)
+                    except Exception as exc:
+                        log_event(self.logger, "worker_sweep_failed",
+                                  message="worker sweep failed: %s" % exc)
                 # 每轮结束写心跳：哪怕本轮一个任务都没派，worker 活着
                 # 这件事本身也要被 server 看到（文件心跳 + 注册表心跳）。
                 self._write_heartbeat()
@@ -1281,16 +1496,27 @@ class Worker:
                         break
                     time.sleep(0.2)
         finally:
-            if self._running_task_id is not None:
-                task_id = self._running_task_id
-                if self._running_proc is not None:
-                    self._kill_proc_group(self._running_proc)
+            hb_stop.set()
+            # 优雅排空：先等正在跑的任务自然结束，超时后才杀。
+            # 串行模式下 run_all 返回时任务已跑完，这里立刻通过；
+            # 并行模式下停机信号退出时后台线程可能还在跑，先等
+            # drain_timeout 秒，超时后按下面杀残留。
+            if drain_timeout is not None:
+                self._wait_for_running(timeout=drain_timeout)
+            for task_id, entry in self._running_snapshot().items():
+                if entry.get("proc") is not None:
+                    self._kill_proc_group(entry["proc"])
                     try:
-                        self._running_proc.wait(timeout=3)
+                        entry["proc"].wait(timeout=3)
                     except Exception:
                         pass
-                self._running_proc = None
-                self._running_task_id = None
+                if entry.get("container") is not None and \
+                        hasattr(self.executor, "kill_container"):
+                    try:
+                        self.executor.kill_container(entry["container"])
+                    except Exception:
+                        pass
+                self._running_clear(task_id)
                 self._set_status(task_id, "interrupted")
         return all_results
 
@@ -1360,8 +1586,9 @@ def main(argv=None):
                         help="多 worker 池身份（默认：主机名-进程号；"
                              "生产部署应配稳定 id，否则重启后旧行变 stale）")
     parser.add_argument("--worker-capacity", type=int, default=1,
-                        help="worker 并发容量声明（注册表里展示用，"
-                             "本轮 run_loop 仍是串行一轮一任务）")
+                        help="worker 并发上限：同一 worker 内最多同时跑"
+                             "几个任务（默认 1，即旧串行行为；>1 时用"
+                             "线程池并行，认领仍是 CAS，成本写库持锁）")
     parser.add_argument("--docker-image", default="codex-app-server:latest")
     parser.add_argument("--docker-cpus", default="2")
     parser.add_argument("--docker-memory", default="2g")
@@ -1386,6 +1613,15 @@ def main(argv=None):
     parser.add_argument("--loop", action="store_true",
                         help="常驻循环，每 --interval 秒跑一轮")
     parser.add_argument("--interval", type=int, default=300)
+    parser.add_argument("--drain-timeout", type=float, default=60,
+                        help="停机优雅排空等待上限（秒）：SIGTERM/SIGINT"
+                             "后先等在跑任务自然结束，超时后才杀残留"
+                             "并标 interrupted；默认 60 秒（systemd"
+                             "默认 90 秒停机时限内能完成）")
+    parser.add_argument("--sweep-threshold", type=float, default=300,
+                        help="自动陈旧回收阈值（秒）：running 任务心跳"
+                             "超过该秒数未更新即被回收；活任务每 30"
+                             "秒刷心跳，不会被误杀")
     parser.add_argument("--max-seconds", type=float, default=None)
     args = parser.parse_args(argv)
     setup_logging()
@@ -1426,7 +1662,9 @@ def main(argv=None):
     try:
         if args.loop:
             results = worker.run_loop(interval=args.interval,
-                                      max_seconds=args.max_seconds)
+                                      max_seconds=args.max_seconds,
+                                      drain_timeout=args.drain_timeout,
+                                      sweep_threshold=args.sweep_threshold)
         else:
             results = worker.run_all(max_tasks=args.max_tasks)
     finally: