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
This commit is contained in:
host committed 2026-10-27 11:06:40 +08:00
1 parent c189df4a20
commit ad014393ee
3 files changed
+474 -26

No files matched your search

+22 -8
View File
@@ -913,14 +913,28 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
- 准入:建任务、会话分支、重试在权限校验通过后、写库之前查 - 准入:建任务、会话分支、重试在权限校验通过后、写库之前查
配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、 配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、
quota_usd_micros、used_usd_micros),不建任务、不改状态。 quota_usd_micros、used_usd_micros),不建任务、不改状态。
老用户与未设配额的用户默认不限额,行为与过去一致。注意 老用户与未设配额的用户默认不限额,行为与过去一致。重试按
边界:只做准入校验,不做跑中掐断——已在跑的回合、后续消 操作人查配额、费用仍记到原任务创建者名下(建任务与分支的
息触发的新一轮中途超限不会被停掉(worker 侧没有按金额中途 操作人就是创建者,没有这个错位)。系统任务(GitHub webhook
停转的钩子);重试按操作人查配额、费用仍记到原任务创建者 的 /fix-review,created_by 是 'system')不记到任何用户名下、
名下(建任务与分支的操作人就是创建者,没有这个错位)。 不占配额;要把它算到项目所有者头上需另做归属。
系统任务(GitHub webhook 的 /fix-review,created_by - 执行时加固(worker 侧,`Worker._quota_exceeded`,口径与准入
是 'system')不记到任何用户名下、不占配额;要把它算到项 一致:当月已记账 + 本轮已累计用量,按价目表换算):认领后
目所有者头上需另做归属。 开跑前再查一次(排队期间配额被花完/调低时不跑,任务标
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、 - 接口:`GET /users/me/quota`(自查,返回 month、quota、
used、remaining,不限额时 quota 与 remaining 为 null); used、remaining,不限额时 quota 与 remaining 为 null);
`GET /admin/quotas?user_id=` 与 `POST /admin/quotas` `GET /admin/quotas?user_id=` 与 `POST /admin/quotas`
+248
View File
@@ -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()
+204 -18
View File
@@ -121,6 +121,7 @@ Docker 下 repos_dir 必须放在 extra-mounts 已挂载的目录下,
import argparse import argparse
import json import json
import os import os
from datetime import datetime, timezone
import re import re
import shutil import shutil
import signal import signal
@@ -746,6 +747,73 @@ class Worker:
return budget_exceeded(self.usage_turns, self.max_turns, return budget_exceeded(self.usage_turns, self.max_turns,
self.usage_cost_usd, self.max_cost_usd) 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): def _ready_tasks(self):
"""当前就绪任务列表(pick_next 的批量版)。 """当前就绪任务列表(pick_next 的批量版)。
@@ -977,14 +1045,20 @@ class Worker:
pass pass
def _execute_via_process(self, task, kind="dispatch", 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/ 返回与同步路径同款 result dict(status/stdout/stderr/
host_status/returncode)。取消时返回 status='cancelled' host_status/returncode)。取消时返回 status='cancelled'
的标记,调用方据此保持 interrupted;deadline_seconds 的标记,调用方据此保持 interrupted;deadline_seconds
到期时杀进程组并返回 status='timeout' 的标记,调用方 到期时杀进程组并返回 status='timeout' 的标记,调用方
据此置 failed(超时是失败,不是取消)。 据此置 failed(超时是失败,不是取消)。quota_check 不为
空时每 30 秒查一次配额(轮询本身 50 毫秒一圈,逐圈查会
把库打爆),超限时同样杀进程组并返回 status=
'quota_exceeded' 的内部标记——调用方据此走配额中断收尾,
绝不能把这个标记原样写进 tasks.status(不在合法状态集
里)。Docker/同步路径没有这个轮询钩子,只走边界检查。
""" """
task_id = task["id"] task_id = task["id"]
if self.process_factory is not None: if self.process_factory is not None:
@@ -1002,7 +1076,9 @@ class Worker:
pass pass
cancelled = False cancelled = False
timed_out = False timed_out = False
quota_hit = False
started = time.time() started = time.time()
last_quota_check = 0.0
try: try:
while proc.poll() is None: while proc.poll() is None:
if task_id in self._cancel_requested or self._loop_stop: if task_id in self._cancel_requested or self._loop_stop:
@@ -1014,6 +1090,17 @@ class Worker:
timed_out = True timed_out = True
self._kill_proc_group(proc) self._kill_proc_group(proc)
break 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) time.sleep(0.05)
try: try:
out, err = proc.communicate(timeout=5) out, err = proc.communicate(timeout=5)
@@ -1030,6 +1117,11 @@ class Worker:
"status": "cancelled", "stdout": out or "", "status": "cancelled", "stdout": out or "",
"stderr": err or "", "returncode": proc.returncode, "stderr": err or "", "returncode": proc.returncode,
"host_status": None} "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: if timed_out:
return {"executor": "local", "task_id": task_id, return {"executor": "local", "task_id": task_id,
"status": "timeout", "stdout": out or "", "status": "timeout", "stdout": out or "",
@@ -1483,7 +1575,7 @@ class Worker:
} }
return (True, tenant_mounts, tenant_network, proxy_env, None) 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,再跑首轮。 """分支任务执行:host.py fork 建分支 thread,再跑首轮。
返回 (result, error),语义同 _run_attempt。fork 本身不 返回 (result, error),语义同 _run_attempt。fork 本身不
跑模型;首轮走 _run_attempt(dispatch 撞已存在 → run 跑模型;首轮走 _run_attempt(dispatch 撞已存在 → run
@@ -1532,7 +1624,7 @@ class Worker:
err = err.strip() err = err.strip()
return (fork_res, return (fork_res,
err[-1000:] if err else "fork failed") 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): def _claim_next_followup(self, task_id):
"""取下一条后续消息:direct 优先、队内按创建时间先来 """取下一条后续消息:direct 优先、队内按创建时间先来
@@ -1595,7 +1687,7 @@ class Worker:
except Exception: except Exception:
pass pass
def _run_attempt(self, task): def _run_attempt(self, task, quota_check=None):
"""单次执行尝试(不含认领、key 门、重试、成本与最终审计)。 """单次执行尝试(不含认领、key 门、重试、成本与最终审计)。
返回 (result, error):result 是执行器结果 dict,error 是 返回 (result, error):result 是执行器结果 dict,error 是
@@ -1662,15 +1754,18 @@ class Worker:
result = self._execute_via_process( result = self._execute_via_process(
task, kind=popen_kind, task, kind=popen_kind,
deadline_seconds=self.task_timeout, 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 "") + \ combined = (result.get("stdout", "") or "") + \
(result.get("stderr", "") or "") (result.get("stderr", "") or "")
if not goal_mode and result["status"] not in ( if not goal_mode and result["status"] not in (
"completed", "cancelled", "timeout") and \ "completed", "cancelled", "timeout",
"quota_exceeded") and \
_ALREADY_EXISTS_MARKER in combined: _ALREADY_EXISTS_MARKER in combined:
result = self._execute_via_process( result = self._execute_via_process(
task, kind="run", task, kind="run",
deadline_seconds=self.task_timeout) deadline_seconds=self.task_timeout,
quota_check=quota_check)
else: else:
if is_docker: if is_docker:
(ok, tenant_mounts, tenant_network, proxy_env, (ok, tenant_mounts, tenant_network, proxy_env,
@@ -1935,6 +2030,30 @@ class Worker:
"stdout": "", "error": repo_error, "events": [], "stdout": "", "error": repo_error, "events": [],
"cost_record_id": None} "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 = [] events = []
fallback_count = 0 fallback_count = 0
result = None result = None
@@ -1963,18 +2082,47 @@ class Worker:
# 各次尝试的真实用量都要算进成本(失败尝试花的钱也是钱): # 各次尝试的真实用量都要算进成本(失败尝试花的钱也是钱):
# 主尝试与 fallback 尝试的结果都先累加,最后只记一条汇总。 # 主尝试与 fallback 尝试的结果都先累加,最后只记一条汇总。
usage_total = None 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);首轮失败 # 分支任务:先 fork 再跑首轮(见 _run_fork);首轮失败
# 不自动重试/fallback,失败后走重试入口重跑首轮(已存在 # 不自动重试/fallback,失败后走重试入口重跑首轮(已存在
# 续跑),不重复 fork。range(0) 让下面的尝试循环跳过, # 续跑),不重复 fork。range(0) 让下面的尝试循环跳过,
# 直接进统一的收尾(超时/取消/成本/审计/通知)逻辑。 # 直接进统一的收尾(超时/取消/成本/审计/通知)逻辑。
if task.get("forked_from"): 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, usage_total = sum_usage(usage_total,
(result or {}).get("usage")) (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") for attempt in range(0 if task.get("forked_from")
else self.max_retries + 1): 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")) 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 \ if self._was_cancelled(task_id) or \
result.get("status") == "cancelled": result.get("status") == "cancelled":
break break
@@ -2011,9 +2159,19 @@ class Worker:
actual_model = fb actual_model = fb
task = dict(task) task = dict(task)
task["model"] = fb 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, usage_total = sum_usage(usage_total,
(result or {}).get("usage")) (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 break
if self.router is not None and host_status_try is not None \ if self.router is not None and host_status_try is not None \
and is_fallbackable(host_status_try) and \ and is_fallbackable(host_status_try) and \
@@ -2045,17 +2203,33 @@ class Worker:
# 剩余 pending 留给重试入口(重试走已存在续跑,跑完继续 # 剩余 pending 留给重试入口(重试走已存在续跑,跑完继续
# 排空剩余消息);取消时该条回 pending,任务走中断收尾。 # 排空剩余消息);取消时该条回 pending,任务走中断收尾。
if (result is not None and result.get("status") == "completed" if (result is not None and result.get("status") == "completed"
and not quota_stop
and not task.get("verify_command") and not task.get("verify_command")
and not task.get("goal_mode")): and not task.get("goal_mode")):
while True: while True:
fu = self._claim_next_followup(task_id) fu = self._claim_next_followup(task_id)
if fu is None: if fu is None:
break 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 = dict(task)
fu_task["prompt"] = fu["body"] 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 = sum_usage(
usage_total, (fu_result or {}).get("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 \ if self._was_cancelled(task_id) or \
(fu_result or {}).get("status") == "cancelled": (fu_result or {}).get("status") == "cancelled":
self._restore_followup(fu["id"]) self._restore_followup(fu["id"])
@@ -2097,7 +2271,16 @@ class Worker:
stdout = result.get("stdout", "") or "" stdout = result.get("stdout", "") or ""
host_status = result.get("host_status") host_status = result.get("host_status")
final = result["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: if fallback_model is not None:
# 真 fallback 后第二次也失败:错误文本必须同时写清 # 真 fallback 后第二次也失败:错误文本必须同时写清
# 两次原因,不能只留第二次的(超时转换已在前面把 # 两次原因,不能只留第二次的(超时转换已在前面把
@@ -2175,11 +2358,13 @@ class Worker:
source="estimated", fallback_count=fallback_count) source="estimated", fallback_count=fallback_count)
if final_written: if final_written:
self._audit(task["project_id"], self._audit(task["project_id"],
"worker_completed" if final == "completed" "worker_quota_stop" if quota_stop else
else "worker_failed", task_id) ("worker_completed" if final == "completed"
else "worker_failed"), task_id)
log_event(self.logger, log_event(self.logger,
"worker_completed" if final == "completed" "worker_quota_stop" if quota_stop else
else "worker_failed", ("worker_completed" if final == "completed"
else "worker_failed"),
task_id=task_id, project_id=task["project_id"], task_id=task_id, project_id=task["project_id"],
status=final) status=final)
if result.get("goal_met") is False: if result.get("goal_met") is False:
@@ -2196,10 +2381,11 @@ class Worker:
# 补一条 worker_completed/failed 审计误导排查。 # 补一条 worker_completed/failed 审计误导排查。
final = "interrupted" final = "interrupted"
self._audit(task["project_id"], "worker_interrupted", task_id) self._audit(task["project_id"], "worker_interrupted", task_id)
if resume_message: if resume_message and not quota_stop:
# 本轮已跑完(无论成败):续跑消息已消费,清掉;进程 # 本轮已跑完(无论成败):续跑消息已消费,清掉;进程
# 中途被杀时清不到,下次认领会重跑同一条消息——可接受 # 中途被杀时清不到,下次认领会重跑同一条消息——可接受
# 的续跑语义,不会静默丢消息。清失败不挡结果返回。 # 的续跑语义,不会静默丢消息。清失败不挡结果返回。
# 配额中断时不消费:顶上来后续跑还要说这句话。
try: try:
conn = _db_connect(self.db_path) conn = _db_connect(self.db_path)
try: try: