From ad014393eeaa08f6bbf1f6c06771e0e094bb28ec Mon Sep 17 00:00:00 2001 From: host Date: Tue, 27 Oct 2026 11:06:40 +0800 Subject: [PATCH] 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/README.md | 30 +++- cloud/tests/test_worker_quota.py | 248 +++++++++++++++++++++++++++++++ cloud/worker.py | 222 ++++++++++++++++++++++++--- 3 files changed, 474 insertions(+), 26 deletions(-) create mode 100644 cloud/tests/test_worker_quota.py diff --git a/cloud/README.md b/cloud/README.md index f3d65d6..cfe629b 100644 --- a/cloud/README.md +++ b/cloud/README.md @@ -913,14 +913,28 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳 - 准入：建任务、会话分支、重试在权限校验通过后、写库之前查 配额，本月用量 >= 配额时回 402 `quota_exceeded`（附 month、 quota_usd_micros、used_usd_micros），不建任务、不改状态。 - 老用户与未设配额的用户默认不限额，行为与过去一致。注意 - 边界：只做准入校验，不做跑中掐断——已在跑的回合、后续消 - 息触发的新一轮中途超限不会被停掉（worker 侧没有按金额中途 - 停转的钩子）；重试按操作人查配额、费用仍记到原任务创建者 - 名下（建任务与分支的操作人就是创建者，没有这个错位）。 - 系统任务（GitHub webhook 的 /fix-review，created_by - 是 'system'）不记到任何用户名下、不占配额；要把它算到项 - 目所有者头上需另做归属。 + 老用户与未设配额的用户默认不限额，行为与过去一致。重试按 + 操作人查配额、费用仍记到原任务创建者名下（建任务与分支的 + 操作人就是创建者，没有这个错位）。系统任务（GitHub webhook + 的 /fix-review，created_by 是 'system'）不记到任何用户名下、 + 不占配额；要把它算到项目所有者头上需另做归属。 +- 执行时加固（worker 侧，`Worker._quota_exceeded`，口径与准入 + 一致：当月已记账 + 本轮已累计用量，按价目表换算）：认领后 + 开跑前再查一次（排队期间配额被花完/调低时不跑，任务标 + interrupted + 配额原因、不记成本、续跑消息不消费）；重试与 + fallback 换模型前各查一次；后续消息每轮开跑前查一次（超限 + 时该条回 pending、任务 interrupted，不再排空剩余；建消息时 + 不查，排队本身不花钱）；Local popen 路径在轮询里每 30 秒查 + 一次，超限时杀进程组（goal 长循环也一样，杀的是整个子进程 + 组）。配额中断的任务一律 interrupted（可续跑：顶上来后重试 + /续跑继续），错误文本带 `quota_exceeded`，已花成本照记 + actual，不发终态推送（与取消同款）。诚实边界：Docker/同步 + 路径没有轮询钩子，只有边界检查，单轮内的超支要到轮次结束 + 才停；单轮本身花的钱无法中途按 token 流式截断（中转用量只 + 在回合结束才回传），所以“掐断”粒度是轮次边界 + 长轮询中的 + 30 秒检查，不是 token 级；并发下两个任务同时查到未超限再同 + 时开跑会有少量超支（与限流器同款 best-effort，非精确配额）； + PR 任务（_run_pr_jobs）暂不查配额。 - 接口：`GET /users/me/quota`（自查，返回 month、quota、 used、remaining，不限额时 quota 与 remaining 为 null）； `GET /admin/quotas?user_id=` 与 `POST /admin/quotas` diff --git a/cloud/tests/test_worker_quota.py b/cloud/tests/test_worker_quota.py new file mode 100644 index 0000000..1fb6017 --- /dev/null +++ b/cloud/tests/test_worker_quota.py @@ -0,0 +1,248 @@ +"""cloud worker 配额加固单测：认领门、后续消息执行时查配额、轮询掐断。 + +运行（在 ~/workspace/codex-app-server 下）： + python3 -m unittest discover -s cloud/tests -p "test_worker_quota.py" -v + +说明：只用注入的假 runner（假 CLI）与真 sleep 子进程，不跑真实 +模型回合、不联网、不登录任何服务；env 映射也是注入的普通 +dict，里面的值只是占位用的假名，绝不是真实密钥。 + +价目（cloud/cost.py，gpt-6.1-sol：input $2.00 / cached $0.10 / +output $10.00 每百万 token）：input 10000 + output 1000 = +20000 + 10000 = 30000 微美元。配额 100000 微美元（$0.10）。 +""" + +import json +import os +import sqlite3 +import subprocess +import sys +import tempfile +import time +import unittest +from pathlib import Path +from types import SimpleNamespace + +try: + from cloud.cost import CostStore + 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.cost import CostStore + from cloud.server import init_db + from cloud.worker import Worker + +QUOTA = 100000 + + +def _insert_task(db_path, task_id, created_by="tester", status="queued", + 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(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)", + (task_id, project_id, "do work", status, created_by, now, now), + ) + conn.commit() + finally: + conn.close() + + +def _insert_quota(db_path, user_id, micros): + conn = sqlite3.connect(db_path) + try: + conn.execute( + "INSERT INTO user_quotas(user_id, quota_usd_micros, updated_at)" + " VALUES(?, ?, ?)", + (user_id, micros, time.time()), + ) + conn.commit() + finally: + conn.close() + + +def _insert_donor_cost(db_path, micros): + """给 tester 记一笔已花成本：另起一个已完成任务并记 actual 行。""" + _insert_task(db_path, "donor", status="completed") + store = CostStore(db_path) + try: + store.record("donor", "p1", "gpt-6.1-sol", input_tokens=1, + output_tokens=0, usd_micros=micros, source="actual") + finally: + try: + store._conn.close() + except Exception: + pass + + +def _insert_followup(db_path, fu_id, task_id, body="follow up"): + now = time.time() + conn = sqlite3.connect(db_path) + try: + conn.execute( + "INSERT INTO task_followups(id, task_id, project_id, body," + " mode, status, created_by, created_at, updated_at)" + " VALUES(?, ?, 'p1', ?, 'queued', 'pending', 'tester', ?, ?)", + (fu_id, task_id, body, 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 _error_text(db_path, task_id): + conn = sqlite3.connect(db_path) + try: + row = conn.execute( + "SELECT error_text FROM tasks WHERE id = ?", + (task_id,)).fetchone() + finally: + conn.close() + return row[0] if row else None + + +def _followup_status(db_path, fu_id): + conn = sqlite3.connect(db_path) + try: + row = conn.execute( + "SELECT status FROM task_followups WHERE id = ?", + (fu_id,)).fetchone() + finally: + conn.close() + return row[0] if row else None + + +def _cost_rows(db_path, task_id): + conn = sqlite3.connect(db_path) + try: + rows = conn.execute( + "SELECT usd_micros, source FROM cost_records" + " WHERE task_id = ?", (task_id,)).fetchall() + finally: + conn.close() + return rows + + +def _usage_runner(usage, calls): + """假 CLI：dispatch/run 一律成功，stdout JSON 里带 usage。""" + def fake(cmd): + calls.append(list(cmd)) + return SimpleNamespace( + returncode=0, + stdout=json.dumps({"status": "completed", "text": "done", + "usage": usage}), + stderr="") + return fake + + +class WorkerQuotaTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + init_db(self.db_path) + self.env = {"XCM_API_KEY": "test-placeholder"} + self.calls = [] + + def tearDown(self): + self.tmp.cleanup() + + def _worker(self, runner): + return Worker(self.db_path, "/fake/host.py", runner=runner, + env=self.env, model="gpt-6.1-sol", + retry_backoff_seconds=0) + + # -- 认领门 -------------------------------------------------------- + def test_claim_time_quota_stop(self): + _insert_quota(self.db_path, "tester", QUOTA) + _insert_donor_cost(self.db_path, QUOTA) + _insert_task(self.db_path, "T1") + worker = self._worker(_usage_runner( + {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, + self.calls)) + result = worker.run_one() + self.assertEqual(result["status"], "interrupted") + self.assertIn("quota_exceeded", result["error"]) + self.assertEqual(_status(self.db_path, "T1"), "interrupted") + self.assertIn("quota_exceeded", _error_text(self.db_path, "T1")) + # 还没跑：runner 一次没调，也不记成本。 + self.assertEqual(self.calls, []) + self.assertEqual(_cost_rows(self.db_path, "T1"), []) + + # -- 后续消息执行时查配额 ------------------------------------------ + def test_followup_quota_stop(self): + _insert_quota(self.db_path, "tester", QUOTA) + _insert_donor_cost(self.db_path, 90000) + _insert_task(self.db_path, "T1") + _insert_followup(self.db_path, "fu1", "T1") + # 首轮花 30000 微美元：90000 + 30000 = 120000 >= 100000， + # 排空时必须停，该条回 pending，任务 interrupted。 + worker = self._worker(_usage_runner( + {"input_tokens": 10000, "cached_input_tokens": 0, + "output_tokens": 1000, "total_tokens": 11000}, + self.calls)) + result = worker.run_one() + self.assertEqual(result["status"], "interrupted") + self.assertIn("quota_exceeded", result["error"]) + self.assertEqual(_status(self.db_path, "T1"), "interrupted") + self.assertEqual(_followup_status(self.db_path, "fu1"), "pending") + # 首轮只跑了一次（后续消息没跑），已花成本照记 actual。 + self.assertEqual(len(self.calls), 1) + self.assertEqual(_cost_rows(self.db_path, "T1"), + [(30000, "actual")]) + + def test_followup_runs_when_under_quota(self): + _insert_quota(self.db_path, "tester", QUOTA) + _insert_task(self.db_path, "T1") + _insert_followup(self.db_path, "fu1", "T1") + worker = self._worker(_usage_runner( + {"input_tokens": 100, "cached_input_tokens": 0, + "output_tokens": 10, "total_tokens": 110}, + self.calls)) + result = worker.run_one() + self.assertEqual(result["status"], "completed") + self.assertEqual(_status(self.db_path, "T1"), "completed") + self.assertEqual(_followup_status(self.db_path, "fu1"), "done") + self.assertEqual(len(self.calls), 2) + + def test_no_quota_row_is_unlimited(self): + _insert_task(self.db_path, "T1") + worker = self._worker(_usage_runner( + {"input_tokens": 100, "output_tokens": 10, "total_tokens": 110}, + self.calls)) + task = worker.get_task("T1") + self.assertIsNone(worker._quota_exceeded(task, None)) + result = worker.run_one() + self.assertEqual(result["status"], "completed") + + # -- 轮询掐断 ------------------------------------------------------ + def test_poll_loop_quota_kill(self): + _insert_task(self.db_path, "T1") + worker = Worker(self.db_path, "/fake/host.py", + env=self.env, model="gpt-6.1-sol", + process_factory=lambda task: subprocess.Popen( + ["sleep", "30"], start_new_session=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE)) + task = worker.get_task("T1") + result = worker._execute_via_process( + task, kind="dispatch", quota_check=lambda: True) + self.assertEqual(result["status"], "quota_exceeded") + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/worker.py b/cloud/worker.py index 61e994a..dcd54eb 100644 --- a/cloud/worker.py +++ b/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: