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/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: