worker: enforce quotas during execution (claim/retry/fallback/followup/poll kill)

- claim-time gate: queued task whose quota was spent while waiting is
  marked interrupted with a quota reason, no cost row, resume kept
- retry/fallback boundaries and every followup round check DB month
  usage plus this run's accumulated usage (priced via cost table)
- local popen poll loop kills the process group on quota (30s throttle);
  Docker/sync paths get boundary checks only (documented)
- quota stops finalize as interrupted + quota_exceeded error, keep the
  actual cost row, no push, followups restored to pending

cloud/worker.py

@@ -121,6 +121,7 @@ Docker 下 repos_dir 必须放在 extra-mounts 已挂载的目录下，
 import argparse
 import json
 import os
+from datetime import datetime, timezone
 import re
 import shutil
 import signal
@@ -746,6 +747,73 @@ class Worker:
         return budget_exceeded(self.usage_turns, self.max_turns,
                                self.usage_cost_usd, self.max_cost_usd)
 
+    def _quota_exceeded(self, task, usage_total):
+        """配额加固的执行时检查：返回超限时的错误文本，否则 None。
+
+        口径与服务端准入（_quota_status）一致：按任务创建者算，
+        当月（UTC）已记账用量 + 本轮已累计用量（按价目表换算，
+        无价模型按 0 计，与服务端 SUM（usd NULL 不计）一致）。
+        无配额行 = 不限额。查询本身异常时按不限额处理（fail-open
+        并记一条警告）：配额是第二道网，绝不能因查询抖动把全部
+        任务挡在门外；服务端建任务/分支/重试时的准入仍是第一道。
+        """
+        owner = (task or {}).get("created_by")
+        if not owner:
+            return None
+        try:
+            month = datetime.now(timezone.utc).strftime("%Y-%m")
+            conn = _db_connect(self.db_path)
+            try:
+                conn.execute(
+                    "CREATE TABLE IF NOT EXISTS user_quotas("
+                    "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT,"
+                    " updated_at DOUBLE PRECISION)")
+                conn.commit()
+                row = conn.execute(
+                    "SELECT quota_usd_micros FROM user_quotas"
+                    " WHERE user_id = ?",
+                    (owner,)).fetchone()
+                if row is None:
+                    return None
+                quota = row[0]
+                if quota is None:
+                    return None
+                # cost_records 由 CostStore 懒建：全新库上它还不存
+                # 在，直接 JOIN 会报错——先实例化一次建表再查（与
+                # 服务端 _quota_status 同款做法），否则 quota=0
+                #（一分钱也不能花）在全新库上会被误放行。
+                _cost_store = CostStore(self.db_path)
+                try:
+                    _cost_store._conn.close()
+                except Exception:
+                    pass
+                used_row = conn.execute(
+                    "SELECT COALESCE(SUM(c.usd_micros), 0)"
+                    " FROM cost_records c JOIN tasks t"
+                    " ON c.task_id = t.id"
+                    " WHERE t.created_by = ? AND c.day LIKE ?",
+                    (owner, month + "%")).fetchone()
+                used = used_row[0] if used_row else 0
+            finally:
+                conn.close()
+            if usage_total:
+                model = (task or {}).get("model") or self.model
+                micros = price_usd_micros(
+                    model, usage_total.get("input_tokens", 0),
+                    usage_total.get("cached_input_tokens", 0),
+                    usage_total.get("output_tokens", 0))
+                used += micros or 0
+            if used >= quota:
+                return ("quota_exceeded: 本月配额已用完（已用 %d / "
+                        "配额 %d 微美元，%s）" % (used, quota, month))
+            return None
+        except Exception as exc:
+            try:
+                self.logger.warning("quota check failed: %s", exc)
+            except Exception:
+                pass
+            return None
+
     def _ready_tasks(self):
         """当前就绪任务列表（pick_next 的批量版）。
 
@@ -977,14 +1045,20 @@ class Worker:
             pass
 
     def _execute_via_process(self, task, kind="dispatch",
-                             deadline_seconds=None, seed_host_dir=None):
+                             deadline_seconds=None, seed_host_dir=None,
+                             quota_check=None):
         """进程模式执行一轮：启动、轮询、取消时杀进程组。
 
         返回与同步路径同款 result dict（status/stdout/stderr/
         host_status/returncode）。取消时返回 status='cancelled'
         的标记，调用方据此保持 interrupted；deadline_seconds
         到期时杀进程组并返回 status='timeout' 的标记，调用方
-        据此置 failed（超时是失败，不是取消）。
+        据此置 failed（超时是失败，不是取消）。quota_check 不为
+        空时每 30 秒查一次配额（轮询本身 50 毫秒一圈，逐圈查会
+        把库打爆），超限时同样杀进程组并返回 status=
+        'quota_exceeded' 的内部标记——调用方据此走配额中断收尾，
+        绝不能把这个标记原样写进 tasks.status（不在合法状态集
+        里）。Docker/同步路径没有这个轮询钩子，只走边界检查。
         """
         task_id = task["id"]
         if self.process_factory is not None:
@@ -1002,7 +1076,9 @@ class Worker:
             pass
         cancelled = False
         timed_out = False
+        quota_hit = False
         started = time.time()
+        last_quota_check = 0.0
         try:
             while proc.poll() is None:
                 if task_id in self._cancel_requested or self._loop_stop:
@@ -1014,6 +1090,17 @@ class Worker:
                     timed_out = True
                     self._kill_proc_group(proc)
                     break
+                if quota_check is not None and \
+                        time.time() - last_quota_check >= 30:
+                    last_quota_check = time.time()
+                    try:
+                        over = quota_check()
+                    except Exception:
+                        over = False
+                    if over:
+                        quota_hit = True
+                        self._kill_proc_group(proc)
+                        break
                 time.sleep(0.05)
             try:
                 out, err = proc.communicate(timeout=5)
@@ -1030,6 +1117,11 @@ class Worker:
                     "status": "cancelled", "stdout": out or "",
                     "stderr": err or "", "returncode": proc.returncode,
                     "host_status": None}
+        if quota_hit:
+            return {"executor": "local", "task_id": task_id,
+                    "status": "quota_exceeded", "stdout": out or "",
+                    "stderr": err or "", "returncode": proc.returncode,
+                    "host_status": None}
         if timed_out:
             return {"executor": "local", "task_id": task_id,
                     "status": "timeout", "stdout": out or "",
@@ -1483,7 +1575,7 @@ class Worker:
             }
         return (True, tenant_mounts, tenant_network, proxy_env, None)
 
-    def _run_fork(self, task):
+    def _run_fork(self, task, quota_check=None):
         """分支任务执行：host.py fork 建分支 thread，再跑首轮。
         返回 (result, error)，语义同 _run_attempt。fork 本身不
         跑模型；首轮走 _run_attempt（dispatch 撞已存在 → run
@@ -1532,7 +1624,7 @@ class Worker:
             err = err.strip()
             return (fork_res,
                     err[-1000:] if err else "fork failed")
-        return self._run_attempt(task)
+        return self._run_attempt(task, quota_check=quota_check)
 
     def _claim_next_followup(self, task_id):
         """取下一条后续消息：direct 优先、队内按创建时间先来
@@ -1595,7 +1687,7 @@ class Worker:
         except Exception:
             pass
 
-    def _run_attempt(self, task):
+    def _run_attempt(self, task, quota_check=None):
         """单次执行尝试（不含认领、key 门、重试、成本与最终审计）。
 
         返回 (result, error)：result 是执行器结果 dict，error 是
@@ -1662,15 +1754,18 @@ class Worker:
                 result = self._execute_via_process(
                     task, kind=popen_kind,
                     deadline_seconds=self.task_timeout,
-                    seed_host_dir=seed_host_dir)
+                    seed_host_dir=seed_host_dir,
+                    quota_check=quota_check)
                 combined = (result.get("stdout", "") or "") + \
                     (result.get("stderr", "") or "")
                 if not goal_mode and result["status"] not in (
-                        "completed", "cancelled", "timeout") and \
+                        "completed", "cancelled", "timeout",
+                        "quota_exceeded") and \
                         _ALREADY_EXISTS_MARKER in combined:
                     result = self._execute_via_process(
                         task, kind="run",
-                        deadline_seconds=self.task_timeout)
+                        deadline_seconds=self.task_timeout,
+                        quota_check=quota_check)
             else:
                 if is_docker:
                     (ok, tenant_mounts, tenant_network, proxy_env,
@@ -1935,6 +2030,30 @@ class Worker:
                     "stdout": "", "error": repo_error, "events": [],
                     "cost_record_id": None}
 
+        # 配额认领门：排队期间配额被花完（或被调低）时，任务不能
+        # 再开跑——直接标 interrupted + 配额原因，不记成本（还没花
+        # 钱），续跑消息不消费，顶上来后重试/续跑可再跑。与 key
+        # 门、仓库门同款提前返回形状，区别是状态为 interrupted
+        #（可续跑）而非 failed。
+        quota_error = self._quota_exceeded(task, None)
+        if quota_error is not None:
+            self._set_status(task_id, "interrupted")
+            self._write_execution_result(
+                task_id, error_text=quota_error, actual_model=self.model,
+                started_at=started_at, finished_at=time.time())
+            self._audit(task["project_id"], "worker_quota_stop",
+                        "%s: %s" % (task_id, quota_error))
+            log_event(self.logger, "worker_quota_stop", task_id=task_id,
+                      project_id=task["project_id"],
+                      status="interrupted")
+            return {"dispatched": True, "task_id": task_id,
+                    "host_name": self.executor._host_name(task)
+                    if hasattr(self.executor, "_host_name") else task_id,
+                    "status": "interrupted",
+                    "executor": self._executor_name(),
+                    "stdout": "", "error": quota_error, "events": [],
+                    "cost_record_id": None}
+
         events = []
         fallback_count = 0
         result = None
@@ -1963,18 +2082,47 @@ class Worker:
         # 各次尝试的真实用量都要算进成本（失败尝试花的钱也是钱）：
         # 主尝试与 fallback 尝试的结果都先累加，最后只记一条汇总。
         usage_total = None
+        # 配额中断状态：任一边界（重试/fallback/后续消息/轮询掐
+        # 断）命中后置位，跳过后续轮次，走统一的配额中断收尾
+        #（interrupted + 配额原因 + 已花成本照记 + 续跑消息保留）。
+        quota_stop = False
+        quota_error = None
+
+        def _quota_hit():
+            # 闭包读当前值：usage_total 与 task 在重试/fallback
+            # 时会被重绑，轮询掐断要按最新值判断。
+            return self._quota_exceeded(task, usage_total) is not None
+
         # 分支任务：先 fork 再跑首轮（见 _run_fork）；首轮失败
         # 不自动重试/fallback，失败后走重试入口重跑首轮（已存在
         # 续跑），不重复 fork。range(0) 让下面的尝试循环跳过，
         # 直接进统一的收尾（超时/取消/成本/审计/通知）逻辑。
         if task.get("forked_from"):
-            result, error = self._run_fork(task)
+            result, error = self._run_fork(task, quota_check=_quota_hit)
             usage_total = sum_usage(usage_total,
                                     (result or {}).get("usage"))
+            if (result or {}).get("status") == "quota_exceeded":
+                quota_error = self._quota_exceeded(task, usage_total) or \
+                    "quota_exceeded: 本月配额已用完"
+                quota_stop = True
         for attempt in range(0 if task.get("forked_from")
                              else self.max_retries + 1):
-            result, error = self._run_attempt(task)
+            if quota_stop:
+                break
+            if attempt > 0:
+                # 重试边界：上一轮已花钱，跑下一轮前再查一次。
+                quota_error = self._quota_exceeded(task, usage_total)
+                if quota_error is not None:
+                    quota_stop = True
+                    break
+            result, error = self._run_attempt(task, quota_check=_quota_hit)
             usage_total = sum_usage(usage_total, (result or {}).get("usage"))
+            if (result or {}).get("status") == "quota_exceeded":
+                # 轮询掐断：进程已被杀，配额原因以现查为准。
+                quota_error = self._quota_exceeded(task, usage_total) or \
+                    "quota_exceeded: 本月配额已用完"
+                quota_stop = True
+                break
             if self._was_cancelled(task_id) or \
                     result.get("status") == "cancelled":
                 break
@@ -2011,9 +2159,19 @@ class Worker:
                 actual_model = fb
                 task = dict(task)
                 task["model"] = fb
-                result, error = self._run_attempt(task)
+                # fallback 边界：换模型再跑一轮前先查配额，超限
+                # 时不烧第二份钱，直接走配额中断收尾。
+                quota_error = self._quota_exceeded(task, usage_total)
+                if quota_error is not None:
+                    quota_stop = True
+                    break
+                result, error = self._run_attempt(task, quota_check=_quota_hit)
                 usage_total = sum_usage(usage_total,
                                         (result or {}).get("usage"))
+                if (result or {}).get("status") == "quota_exceeded":
+                    quota_error = self._quota_exceeded(task, usage_total) or \
+                        "quota_exceeded: 本月配额已用完"
+                    quota_stop = True
                 break
             if self.router is not None and host_status_try is not None \
                     and is_fallbackable(host_status_try) and \
@@ -2045,17 +2203,33 @@ class Worker:
         # 剩余 pending 留给重试入口（重试走已存在续跑，跑完继续
         # 排空剩余消息）；取消时该条回 pending，任务走中断收尾。
         if (result is not None and result.get("status") == "completed"
+                and not quota_stop
                 and not task.get("verify_command")
                 and not task.get("goal_mode")):
             while True:
                 fu = self._claim_next_followup(task_id)
                 if fu is None:
                     break
+                # 后续消息执行时查配额：超限时该条回 pending（留
+                # 给重试入口），任务走配额中断收尾，不再排空剩余；
+                # 建消息时不查（排队本身不花钱，顶上来后可跑）。
+                quota_error = self._quota_exceeded(task, usage_total)
+                if quota_error is not None:
+                    self._restore_followup(fu["id"])
+                    quota_stop = True
+                    break
                 fu_task = dict(task)
                 fu_task["prompt"] = fu["body"]
-                fu_result, fu_error = self._run_attempt(fu_task)
+                fu_result, fu_error = self._run_attempt(
+                    fu_task, quota_check=_quota_hit)
                 usage_total = sum_usage(
                     usage_total, (fu_result or {}).get("usage"))
+                if (fu_result or {}).get("status") == "quota_exceeded":
+                    self._restore_followup(fu["id"])
+                    quota_error = self._quota_exceeded(task, usage_total) or \
+                        "quota_exceeded: 本月配额已用完"
+                    quota_stop = True
+                    break
                 if self._was_cancelled(task_id) or \
                         (fu_result or {}).get("status") == "cancelled":
                     self._restore_followup(fu["id"])
@@ -2097,7 +2271,16 @@ class Worker:
         stdout = result.get("stdout", "") or ""
         host_status = result.get("host_status")
         final = result["status"]
-        if final != "completed":
+        if quota_stop:
+            # 配额中断收尾：状态一律 interrupted + 配额原因，不跑
+            # 失败原因重算（否则 fallback 拼接会把配额原因盖掉），
+            # 已花成本仍按 usage_total 照记（下面成本块），续跑
+            # 消息保留（下面清消息时跳过），不发终态推送（与取消
+            # 同款：interrupted 不在推送事件集里）。quota_exceeded
+            # 这个内部标记绝不能原样写进 tasks.status。
+            final = "interrupted"
+            error = quota_error
+        if final != "completed" and not quota_stop:
             if fallback_model is not None:
                 # 真 fallback 后第二次也失败：错误文本必须同时写清
                 # 两次原因，不能只留第二次的（超时转换已在前面把
@@ -2175,11 +2358,13 @@ class Worker:
                     source="estimated", fallback_count=fallback_count)
         if final_written:
             self._audit(task["project_id"],
-                        "worker_completed" if final == "completed"
-                        else "worker_failed", task_id)
+                        "worker_quota_stop" if quota_stop else
+                        ("worker_completed" if final == "completed"
+                         else "worker_failed"), task_id)
             log_event(self.logger,
-                      "worker_completed" if final == "completed"
-                      else "worker_failed",
+                      "worker_quota_stop" if quota_stop else
+                      ("worker_completed" if final == "completed"
+                       else "worker_failed"),
                       task_id=task_id, project_id=task["project_id"],
                       status=final)
             if result.get("goal_met") is False:
@@ -2196,10 +2381,11 @@ class Worker:
             # 补一条 worker_completed/failed 审计误导排查。
             final = "interrupted"
             self._audit(task["project_id"], "worker_interrupted", task_id)
-        if resume_message:
+        if resume_message and not quota_stop:
             # 本轮已跑完（无论成败）：续跑消息已消费，清掉；进程
             # 中途被杀时清不到，下次认领会重跑同一条消息——可接受
             # 的续跑语义，不会静默丢消息。清失败不挡结果返回。
+            # 配额中断时不消费：顶上来后续跑还要说这句话。
             try:
                 conn = _db_connect(self.db_path)
                 try: