feat: 成本预算与失败通知第一批(G1/G2/G5/G6/无价上报)
- G1:修正 _check_quota docstring(Local popen 30 秒轮询掐断、 Docker/同步路径只有边界检查),README 配额节同步;followup 创建不查配额是已声明设计,保持不动。 - G2:notify_failure 对插值 shlex.quote,模板错误记日志不静默 吞,命令返回值/超时/异常记 host-run.log;dispatch 前置预算门 与 goal 预算中断补 budget-exceeded 通知(没配命令时只写日志)。 - G5/G6:重试与续跑的配额检查按任务创建者查(费用记创建者名下), 维护者替他人操作不再用错人的额度;续跑 API 侧补查,超限回 402 且不改状态,不再先 200 排队再悄悄 interrupted。 - 无价上报:配额自查/管理端点返回增加 unpriced_actual_tokens (actual 且 usd_micros 为 NULL,与 CostStore 口径一致),只读 上报、不计价、不拦任务;不给 gpt-5.6-terra 编造价格。 测试:host 68、cloud 760(4 跳过)、deploy 19 全过。
This commit is contained in:
1 parent
c2f82da5e0
commit
e4b5545a24
5 files changed
+304
-38
No files matched your search
+18
-10
@@ -956,18 +956,22 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
|||||||
NULL,SUM 自动跳过),按任务创建者 `created_by` 归属——一个
|
NULL,SUM 自动跳过),按任务创建者 `created_by` 归属——一个
|
||||||
用户在多个项目里花的钱都算进他自己的月度账单;项目维度的
|
用户在多个项目里花的钱都算进他自己的月度账单;项目维度的
|
||||||
拆分由项目成本接口另行提供。
|
拆分由项目成本接口另行提供。
|
||||||
- 准入:建任务、会话分支、重试在权限校验通过后、写库之前查
|
- 准入:建任务、会话分支、重试、续跑在权限校验通过后、写库
|
||||||
配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、
|
之前查配额,本月用量 >= 配额时回 402 `quota_exceeded`(附
|
||||||
quota_usd_micros、used_usd_micros),不建任务、不改状态。
|
month、quota_usd_micros、used_usd_micros),不建任务、不改
|
||||||
老用户与未设配额的用户默认不限额,行为与过去一致。但
|
状态。老用户与未设配额的用户默认不限额,行为与过去一致。但
|
||||||
OAuth 新用户不再是“建号即不限额”:建号(created=True)时
|
OAuth 新用户不再是“建号即不限额”:建号(created=True)时
|
||||||
写一行默认配额(`CLOUD_DEFAULT_QUOTA_USD_MICROS`,默认
|
写一行默认配额(`CLOUD_DEFAULT_QUOTA_USD_MICROS`,默认
|
||||||
10_000_000,即 $10/月;`INSERT OR IGNORE` 幂等,不覆盖管理
|
10_000_000,即 $10/月;`INSERT OR IGNORE` 幂等,不覆盖管理
|
||||||
员后来改过的额度),并记 `quota_default` 审计。重试按
|
员后来改过的额度),并记 `quota_default` 审计。重试与续跑
|
||||||
操作人查配额、费用仍记到原任务创建者名下(建任务与分支的
|
按任务创建者查配额(费用本来就记在创建者名下;维护者替
|
||||||
操作人就是创建者,没有这个错位)。系统任务(GitHub webhook
|
他人重试/续跑时,操作人超限不应拦、创建者超限才应拦——
|
||||||
的 /fix-review,created_by 是 'system')不记到任何用户名下、
|
旧版重试按操作人查是错位,本轮已改)。后续消息
|
||||||
不占配额;要把它算到项目所有者头上需另做归属。
|
(followup)创建不查配额:排队本身不花钱,执行时轮次边界
|
||||||
|
的配额门照常生效,不要把 followup 也改成准入校验。系统
|
||||||
|
任务(GitHub webhook 的 /fix-review,created_by 是 'system')
|
||||||
|
不记到任何用户名下、不占配额;要把它算到项目所有者头上
|
||||||
|
需另做归属。
|
||||||
- 无价模型 token 上限:价目表目前只有两个 sol 模型有价,默认
|
- 无价模型 token 上限:价目表目前只有两个 sol 模型有价,默认
|
||||||
模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL,USD
|
模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL,USD
|
||||||
配额对它们永远不生效(按 0 计)。因此有配额行的用户另受一
|
配额对它们永远不生效(按 0 计)。因此有配额行的用户另受一
|
||||||
@@ -997,7 +1001,11 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
|
|||||||
时开跑会有少量超支(与限流器同款 best-effort,非精确配额);
|
时开跑会有少量超支(与限流器同款 best-effort,非精确配额);
|
||||||
PR 任务(_run_pr_jobs)暂不查配额。
|
PR 任务(_run_pr_jobs)暂不查配额。
|
||||||
- 接口:`GET /users/me/quota`(自查,返回 month、quota、
|
- 接口:`GET /users/me/quota`(自查,返回 month、quota、
|
||||||
used、remaining,不限额时 quota 与 remaining 为 null);
|
used、remaining,以及当月无价实际 token 数
|
||||||
|
`unpriced_actual_tokens`——价目表里没有该模型价格的实际
|
||||||
|
token 只上报、不计价、不拦任务,绝不编造一个价格把它折
|
||||||
|
算进 used;不限额时 quota 与 remaining 为 null);管理端点
|
||||||
|
的配额返回同样含 `unpriced_actual_tokens`;
|
||||||
`GET /admin/quotas?user_id=` 与 `POST /admin/quotas`
|
`GET /admin/quotas?user_id=` 与 `POST /admin/quotas`
|
||||||
(body:user_id、quota_usd_micros,null=不限额,负数 400,
|
(body:user_id、quota_usd_micros,null=不限额,负数 400,
|
||||||
用户不存在 404,改完记 `quota_set` 审计)。管理口径沿用
|
用户不存在 404,改完记 `quota_set` 审计)。管理口径沿用
|
||||||
|
|||||||
+71
-20
@@ -3661,6 +3661,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
在旧线程上开新目标循环;goal_mode 不传时 goal 任务沿用
|
在旧线程上开新目标循环;goal_mode 不传时 goal 任务沿用
|
||||||
原语义续跑。状态回 queued 时上一轮 result/error 清空
|
原语义续跑。状态回 queued 时上一轮 result/error 清空
|
||||||
(与重试一致,否则排队中的任务还挂着旧结果)。
|
(与重试一致,否则排队中的任务还挂着旧结果)。
|
||||||
|
配额:全部校验通过后、写库之前按任务创建者查一次
|
||||||
|
(与重试同口径,费用记创建者名下);超限回 402,不
|
||||||
|
改状态——以前是先回 200 排队、随后 worker 认领门把
|
||||||
|
任务悄悄标 interrupted 且不推送。后续消息创建不查
|
||||||
|
配额(排队不花钱),本接口不是 followup,不要混淆。
|
||||||
并发下用 CAS(WHERE status IN 终端四态),没改到行
|
并发下用 CAS(WHERE status IN 终端四态),没改到行
|
||||||
说明状态已变,返回 409 而不是假装成功。写审计
|
说明状态已变,返回 409 而不是假装成功。写审计
|
||||||
task_continued。
|
task_continued。
|
||||||
@@ -3768,6 +3773,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
sets.append("verify_command = NULL")
|
sets.append("verify_command = NULL")
|
||||||
sets.append("goal_max_rounds = 0")
|
sets.append("goal_max_rounds = 0")
|
||||||
vals.append(task_id)
|
vals.append(task_id)
|
||||||
|
# 续跑也要花模型钱:校验全过后、写库之前按创建者查
|
||||||
|
# 配额,超限回 402 且不改状态(见本函数 docstring)。
|
||||||
|
if self._reject_if_over_quota(
|
||||||
|
user_id, quota_user_id=task.get("created_by") or user_id):
|
||||||
|
return
|
||||||
conn = self._db()
|
conn = self._db()
|
||||||
try:
|
try:
|
||||||
cur = conn.execute(
|
cur = conn.execute(
|
||||||
@@ -3802,6 +3812,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
interrupted / merge-failed(其它状态一律 409:running
|
interrupted / merge-failed(其它状态一律 409:running
|
||||||
本来就在跑,queued 本来就在等,completed/merged 已有
|
本来就在跑,queued 本来就在等,completed/merged 已有
|
||||||
结果,重跑应是新任务而不是覆盖结果)。
|
结果,重跑应是新任务而不是覆盖结果)。
|
||||||
|
配额按任务创建者查(费用记在创建者名下):维护者替
|
||||||
|
他人重试时,操作人超限不应拦、创建者超限才应拦。
|
||||||
成功后状态回 queued,上一轮的 result_text/error_text/
|
成功后状态回 queued,上一轮的 result_text/error_text/
|
||||||
actual_model 清空(否则新一轮的结果会和旧结果混在一起
|
actual_model 清空(否则新一轮的结果会和旧结果混在一起
|
||||||
无法分辨);prompt、模型、优先级、依赖原样保留,依赖
|
无法分辨);prompt、模型、优先级、依赖原样保留,依赖
|
||||||
@@ -3820,7 +3832,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
user_id, task["project_id"], "approve_merge"):
|
user_id, task["project_id"], "approve_merge"):
|
||||||
self._send_json(403, {"error": "forbidden"})
|
self._send_json(403, {"error": "forbidden"})
|
||||||
return
|
return
|
||||||
if self._reject_if_over_quota(user_id):
|
if self._reject_if_over_quota(
|
||||||
|
user_id, quota_user_id=task.get("created_by") or user_id):
|
||||||
return
|
return
|
||||||
if task["status"] not in ("failed", "interrupted", "merge-failed"):
|
if task["status"] not in ("failed", "interrupted", "merge-failed"):
|
||||||
self._send_json(409, {
|
self._send_json(409, {
|
||||||
@@ -3899,20 +3912,23 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
finally:
|
finally:
|
||||||
conn.close()
|
conn.close()
|
||||||
|
|
||||||
def _quota_payload(self, user_id, quota, used, month):
|
def _quota_unpriced_tokens(self, user_id, month=None):
|
||||||
return {"user_id": user_id, "month": month,
|
|
||||||
"quota_usd_micros": quota, "used_usd_micros": used,
|
|
||||||
"remaining_usd_micros": (quota - used)
|
|
||||||
if quota is not None else None}
|
|
||||||
|
|
||||||
def _quota_unpriced_tokens(self, user_id, month):
|
|
||||||
"""当月无价模型 token 数(usd_micros IS NULL 的成本行,
|
"""当月无价模型 token 数(usd_micros IS NULL 的成本行,
|
||||||
input + output 之和;cached 已含在 input 里,不重复计)。
|
input + output 之和;cached 已含在 input 里,不重复计)。
|
||||||
|
|
||||||
与 _quota_status 同款 CostStore 懒建表做法,全新库上
|
合并说明(#86 × #90):两边曾各自定义过一次(#90 的
|
||||||
cost_records 不存在时先建表再查。无配额行时调用方不
|
只要 user_id 且只计 source='actual',#86 的要 month 且
|
||||||
应以此挡人(见 _check_unpriced_token_cap)。
|
计全部 usd IS NULL 行),自动合并后后一个覆盖前一个,
|
||||||
|
上报调用点会 TypeError。这里统一成一个方法,以 main
|
||||||
|
(#86)的口径为准:计全部 usd IS NULL 行(含 estimated
|
||||||
|
行),上报与上限共用同一口径,不再留两套数字;month
|
||||||
|
缺省为当月,兼容两边的调用点。与 _quota_status 同款
|
||||||
|
CostStore 懒建表做法,全新库上 cost_records 不存在时
|
||||||
|
先建表再查。无配额行时调用方不应以此挡人(见
|
||||||
|
_check_unpriced_token_cap)。
|
||||||
"""
|
"""
|
||||||
|
if month is None:
|
||||||
|
month = self._quota_month()
|
||||||
conn = self._db()
|
conn = self._db()
|
||||||
try:
|
try:
|
||||||
_cost_store = CostStore(self.server.db_path)
|
_cost_store = CostStore(self.server.db_path)
|
||||||
@@ -3930,6 +3946,16 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
finally:
|
finally:
|
||||||
conn.close()
|
conn.close()
|
||||||
|
|
||||||
|
def _quota_payload(self, user_id, quota, used, month,
|
||||||
|
unpriced_actual_tokens=None):
|
||||||
|
payload = {"user_id": user_id, "month": month,
|
||||||
|
"quota_usd_micros": quota, "used_usd_micros": used,
|
||||||
|
"remaining_usd_micros": (quota - used)
|
||||||
|
if quota is not None else None}
|
||||||
|
if unpriced_actual_tokens is not None:
|
||||||
|
payload["unpriced_actual_tokens"] = unpriced_actual_tokens
|
||||||
|
return payload
|
||||||
|
|
||||||
def _check_unpriced_token_cap(self, user_id):
|
def _check_unpriced_token_cap(self, user_id):
|
||||||
"""无价 token 上限准入:超限返回 (cap, used, month),否则 None。
|
"""无价 token 上限准入:超限返回 (cap, used, month),否则 None。
|
||||||
|
|
||||||
@@ -3958,16 +3984,34 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
def _check_quota(self, user_id):
|
def _check_quota(self, user_id):
|
||||||
"""配额准入:放行返回 None;超限返回 (quota, used,
|
"""配额准入:放行返回 None;超限返回 (quota, used,
|
||||||
month),调用方一律回 402——这是钱用完了,不是请求
|
month),调用方一律回 402——这是钱用完了,不是请求
|
||||||
太快,所以不用 429。注意:只做建任务/分支/重跑时的
|
太快,所以不用 429。
|
||||||
准入校验;跑起来的回合中途超限不会被掐断(worker
|
|
||||||
侧没有按金额中途停转的钩子,那是下一步的事)。"""
|
准入点:建任务 / 分支 / 重试 / 续跑,均为权限校验
|
||||||
|
通过后、写库之前。后续消息(followup)创建不查:
|
||||||
|
排队本身不花钱,执行时轮次边界的配额门照常生效,
|
||||||
|
不要顺手把 followup 也改成准入校验。
|
||||||
|
|
||||||
|
跑中掐断(与旧 docstring 不同,旧说法已过时):
|
||||||
|
Local popen 路径的 worker 轮询每 30 秒查一次配额,
|
||||||
|
超限杀进程组;Docker / 同步路径没有轮询钩子,只有
|
||||||
|
认领与轮次边界检查,单轮内的超支要到轮次结束才停;
|
||||||
|
单轮本身的钱无法按 token 流式截断(用量只在回合
|
||||||
|
结束才回传),掐断粒度是轮次边界 + 长轮询中的
|
||||||
|
30 秒检查,不是 token 级。"""
|
||||||
quota, used, month = self._quota_status(user_id)
|
quota, used, month = self._quota_status(user_id)
|
||||||
if quota is not None and used >= quota:
|
if quota is not None and used >= quota:
|
||||||
return quota, used, month
|
return quota, used, month
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def _reject_if_over_quota(self, user_id):
|
def _reject_if_over_quota(self, user_id, quota_user_id=None):
|
||||||
blocked = self._check_quota(user_id)
|
"""配额超限时回 402 并返回 True,否则 False。
|
||||||
|
|
||||||
|
quota_user_id:实际查谁的配额。建任务 / 分支时调用
|
||||||
|
方就是创建者,缺省即调用方;重试 / 续跑时费用记在
|
||||||
|
原任务创建者名下,必须传 task 的 created_by,否则
|
||||||
|
维护者替他人重试会用错人的额度(旧 bug:重试按
|
||||||
|
操作人查、费用记创建者名下)。"""
|
||||||
|
blocked = self._check_quota(quota_user_id or user_id)
|
||||||
if blocked is not None:
|
if blocked is not None:
|
||||||
quota, used, month = blocked
|
quota, used, month = blocked
|
||||||
self._send_json(402, {"error": "quota_exceeded",
|
self._send_json(402, {"error": "quota_exceeded",
|
||||||
@@ -3975,7 +4019,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
"quota_usd_micros": quota,
|
"quota_usd_micros": quota,
|
||||||
"used_usd_micros": used})
|
"used_usd_micros": used})
|
||||||
return True
|
return True
|
||||||
capped = self._check_unpriced_token_cap(user_id)
|
capped = self._check_unpriced_token_cap(
|
||||||
|
quota_user_id or user_id)
|
||||||
if capped is not None:
|
if capped is not None:
|
||||||
cap, used, month = capped
|
cap, used, month = capped
|
||||||
self._send_json(402, {"error": "quota_exceeded",
|
self._send_json(402, {"error": "quota_exceeded",
|
||||||
@@ -4011,7 +4056,9 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
def _user_quota_self(self, user_id):
|
def _user_quota_self(self, user_id):
|
||||||
quota, used, month = self._quota_status(user_id)
|
quota, used, month = self._quota_status(user_id)
|
||||||
self._send_json(200,
|
self._send_json(200,
|
||||||
self._quota_payload(user_id, quota, used, month))
|
self._quota_payload(
|
||||||
|
user_id, quota, used, month,
|
||||||
|
self._quota_unpriced_tokens(user_id)))
|
||||||
|
|
||||||
def _admin_quota_get(self, user_id, query):
|
def _admin_quota_get(self, user_id, query):
|
||||||
if not self._is_global_maintainer(user_id):
|
if not self._is_global_maintainer(user_id):
|
||||||
@@ -4026,7 +4073,9 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
return
|
return
|
||||||
quota, used, month = self._quota_status(target)
|
quota, used, month = self._quota_status(target)
|
||||||
self._send_json(200,
|
self._send_json(200,
|
||||||
self._quota_payload(target, quota, used, month))
|
self._quota_payload(
|
||||||
|
target, quota, used, month,
|
||||||
|
self._quota_unpriced_tokens(target)))
|
||||||
|
|
||||||
def _admin_quota_set(self, user_id):
|
def _admin_quota_set(self, user_id):
|
||||||
if not self._is_global_maintainer(user_id):
|
if not self._is_global_maintainer(user_id):
|
||||||
@@ -4075,7 +4124,9 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
"%s:%s" % (target, quota))
|
"%s:%s" % (target, quota))
|
||||||
quota, used, month = self._quota_status(target)
|
quota, used, month = self._quota_status(target)
|
||||||
self._send_json(200,
|
self._send_json(200,
|
||||||
self._quota_payload(target, quota, used, month))
|
self._quota_payload(
|
||||||
|
target, quota, used, month,
|
||||||
|
self._quota_unpriced_tokens(target)))
|
||||||
|
|
||||||
# -- worker 池 -----------------------------------------------------
|
# -- worker 池 -----------------------------------------------------
|
||||||
def _workers_list(self, user_id):
|
def _workers_list(self, user_id):
|
||||||
|
|||||||
@@ -234,6 +234,95 @@ class ServerQuotasTestCase(unittest.TestCase):
|
|||||||
self.assertEqual(s, 200, retried)
|
self.assertEqual(s, 200, retried)
|
||||||
self.assertEqual(retried["status"], "queued")
|
self.assertEqual(retried["status"], "queued")
|
||||||
|
|
||||||
|
def test_retry_uses_creator_quota_not_operator(self):
|
||||||
|
# 重试按任务创建者查配额:费用记创建者名下,操作人
|
||||||
|
# 是谁不应影响准入。other 须是 maintainer 才能过
|
||||||
|
# approve_merge 权限门,否则 403 会掩盖配额判断。
|
||||||
|
self.auth.set_member(self.project_id, self.owner_id,
|
||||||
|
self.other_id, "maintainer")
|
||||||
|
s, task = self._create_task()
|
||||||
|
self.assertEqual(s, 201, task)
|
||||||
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||||
|
2000000, _utc_today())
|
||||||
|
_set_task_status(self.db_path, task["id"], "failed")
|
||||||
|
# 创建者超限:维护者替他重试也应拦。
|
||||||
|
s, _ = self._set_quota(self.owner_id, 2000000)
|
||||||
|
self.assertEqual(s, 200)
|
||||||
|
s, body = _req(self.port, "POST",
|
||||||
|
"/tasks/%s/retry" % task["id"], {},
|
||||||
|
self.other_key)
|
||||||
|
self.assertEqual(s, 402, body)
|
||||||
|
# 创建者不限额、操作人额度为 0:按旧口径(查操作人)
|
||||||
|
# 会拦,按创建者口径应放行。
|
||||||
|
s, _ = self._set_quota(self.owner_id, None)
|
||||||
|
self.assertEqual(s, 200)
|
||||||
|
s, _ = self._set_quota(self.other_id, 0)
|
||||||
|
self.assertEqual(s, 200)
|
||||||
|
s, retried = _req(self.port, "POST",
|
||||||
|
"/tasks/%s/retry" % task["id"], {},
|
||||||
|
self.other_key)
|
||||||
|
self.assertEqual(s, 200, retried)
|
||||||
|
self.assertEqual(retried["status"], "queued")
|
||||||
|
|
||||||
|
def test_continue_blocked_at_creator_quota(self):
|
||||||
|
s, task = self._create_task()
|
||||||
|
self.assertEqual(s, 201, task)
|
||||||
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||||
|
5000000, _utc_today())
|
||||||
|
_set_task_status(self.db_path, task["id"], "completed")
|
||||||
|
s, _ = self._set_quota(self.owner_id, 5000000)
|
||||||
|
self.assertEqual(s, 200)
|
||||||
|
# 以前:先回 200 排队、随后认领门悄悄标 interrupted。
|
||||||
|
# 现在:API 侧直接 402,状态保持 completed 不变。
|
||||||
|
s, body = _req(self.port, "POST",
|
||||||
|
"/tasks/%s/continue" % task["id"],
|
||||||
|
{"message": "again"}, self.owner_key)
|
||||||
|
self.assertEqual(s, 402, body)
|
||||||
|
s, row = _req(self.port, "GET", "/tasks/" + task["id"],
|
||||||
|
key=self.owner_key)
|
||||||
|
self.assertEqual(row["status"], "completed")
|
||||||
|
# 额度调高后放行。
|
||||||
|
s, _ = self._set_quota(self.owner_id, None)
|
||||||
|
self.assertEqual(s, 200)
|
||||||
|
s, continued = _req(self.port, "POST",
|
||||||
|
"/tasks/%s/continue" % task["id"],
|
||||||
|
{"message": "again"}, self.owner_key)
|
||||||
|
self.assertEqual(s, 200, continued)
|
||||||
|
self.assertEqual(continued["status"], "queued")
|
||||||
|
# 续跑同样按创建者查:操作人额度为 0、创建者不限额时放行。
|
||||||
|
self.auth.set_member(self.project_id, self.owner_id,
|
||||||
|
self.other_id, "maintainer")
|
||||||
|
s, task2 = self._create_task()
|
||||||
|
self.assertEqual(s, 201, task2)
|
||||||
|
_set_task_status(self.db_path, task2["id"], "completed")
|
||||||
|
s, _ = self._set_quota(self.other_id, 0)
|
||||||
|
self.assertEqual(s, 200)
|
||||||
|
s, continued2 = _req(self.port, "POST",
|
||||||
|
"/tasks/%s/continue" % task2["id"],
|
||||||
|
{"message": "again"}, self.other_key)
|
||||||
|
self.assertEqual(s, 200, continued2)
|
||||||
|
self.assertEqual(continued2["status"], "queued")
|
||||||
|
|
||||||
|
def test_quota_reports_unpriced_tokens(self):
|
||||||
|
s, task = self._create_task()
|
||||||
|
self.assertEqual(s, 201, task)
|
||||||
|
# 有价 5000000 计入 used;无价 actual 150 只进
|
||||||
|
# unpriced_actual_tokens,不计价、不拦任务。
|
||||||
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||||
|
5000000, _utc_today())
|
||||||
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||||
|
None, _utc_today())
|
||||||
|
s, q = _req(self.port, "GET", "/users/me/quota",
|
||||||
|
key=self.owner_key)
|
||||||
|
self.assertEqual(s, 200, q)
|
||||||
|
self.assertEqual(q["used_usd_micros"], 5000000)
|
||||||
|
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
||||||
|
s, q = _req(self.port, "GET",
|
||||||
|
"/admin/quotas?user_id=" + self.owner_id,
|
||||||
|
key=self.owner_key)
|
||||||
|
self.assertEqual(s, 200, q)
|
||||||
|
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
@@ -48,6 +48,7 @@ import json
|
|||||||
import os
|
import os
|
||||||
import queue
|
import queue
|
||||||
import re
|
import re
|
||||||
|
import shlex
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
import threading
|
import threading
|
||||||
@@ -1375,20 +1376,59 @@ def budget_exceeded(state, cfg):
|
|||||||
def notify_failure(state_path, cfg, task_name, result):
|
def notify_failure(state_path, cfg, task_name, result):
|
||||||
"""失败通知:先写 host-run.log(一定有),再跑用户配置的
|
"""失败通知:先写 host-run.log(一定有),再跑用户配置的
|
||||||
notify_command(本地 shell 命令,只来自用户自己的配置,
|
notify_command(本地 shell 命令,只来自用户自己的配置,
|
||||||
不是外部输入;模板里可用 {task} {status})。"""
|
不是外部输入;模板里可用 {task} {status} {text})。
|
||||||
text = (result or {}).get("text", "") or ""
|
|
||||||
msg = f"FAILED {task_name}: {result.get('status')} {text[:200]}"
|
插进命令串的值一律 shlex.quote:模型文本可能含 shell 元
|
||||||
|
字符(引号、分号、$() 等),不加 quote 就等于把模型输出
|
||||||
|
当命令执行。模板本身有问题(未知占位符、未成对的大括号)
|
||||||
|
时不再静默吞掉:记一条 template error 到 host-run.log 并
|
||||||
|
且不跑这条命令。命令的返回值 / 超时 / 异常同样记日志,
|
||||||
|
否则通知丢了都无从发现(旧行为是异常全吞、无任何记录)。
|
||||||
|
没有配 notify_command 时只写 host-run.log,不外发。"""
|
||||||
|
result = result or {}
|
||||||
|
text = result.get("text", "") or ""
|
||||||
|
status = result.get("status", "") or ""
|
||||||
|
msg = f"FAILED {task_name}: {status} {text[:200]}"
|
||||||
try:
|
try:
|
||||||
log_event(state_path, msg)
|
log_event(state_path, msg)
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
cmd = (cfg or {}).get("notify_command")
|
cmd = (cfg or {}).get("notify_command")
|
||||||
if cmd:
|
if not cmd:
|
||||||
|
return msg
|
||||||
|
try:
|
||||||
|
formatted = cmd.format(task=shlex.quote(str(task_name)),
|
||||||
|
status=shlex.quote(str(status)),
|
||||||
|
text=shlex.quote(str(text[:200])))
|
||||||
|
except (KeyError, IndexError, ValueError) as exc:
|
||||||
|
# 模板写错(例如 {unknown} 或单个 {)时 .format 抛错:
|
||||||
|
# 明确记日志、不跑命令、不让调用方以为通知已发出。
|
||||||
try:
|
try:
|
||||||
subprocess.run(cmd.format(task=task_name,
|
log_event(state_path,
|
||||||
status=result.get("status", ""),
|
"notify command template error: "
|
||||||
text=text[:200]),
|
f"{type(exc).__name__}: {exc}")
|
||||||
shell=True, timeout=30)
|
except Exception:
|
||||||
|
pass
|
||||||
|
return msg
|
||||||
|
try:
|
||||||
|
proc = subprocess.run(formatted, shell=True, timeout=30)
|
||||||
|
try:
|
||||||
|
log_event(state_path,
|
||||||
|
f"notify command rc={proc.returncode} "
|
||||||
|
f"task={task_name}")
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
try:
|
||||||
|
log_event(state_path,
|
||||||
|
f"notify command timeout task={task_name}")
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
except Exception as exc:
|
||||||
|
try:
|
||||||
|
log_event(state_path,
|
||||||
|
"notify command failed: "
|
||||||
|
f"{type(exc).__name__} task={task_name}")
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
return msg
|
return msg
|
||||||
@@ -1969,6 +2009,11 @@ def run_goal_loop(server, state, state_path, task, goal_text, verify,
|
|||||||
if budget_stopped:
|
if budget_stopped:
|
||||||
task["status"] = "interrupted"
|
task["status"] = "interrupted"
|
||||||
task["goal_stopped_reason"] = "budget exceeded: max_daily_turns"
|
task["goal_stopped_reason"] = "budget exceeded: max_daily_turns"
|
||||||
|
# goal 预算中断以前只打印:与 dispatch 前置门同款,
|
||||||
|
# 补一次 budget-exceeded 通知(没配命令时只写日志)。
|
||||||
|
notify_failure(state_path, cfg, task.get("name") or "goal",
|
||||||
|
{"status": "budget-exceeded",
|
||||||
|
"text": "budget exceeded: max_daily_turns"})
|
||||||
else:
|
else:
|
||||||
task["status"] = "goal-met" if ok else "goal-unmet"
|
task["status"] = "goal-met" if ok else "goal-unmet"
|
||||||
if task["status"] == "goal-unmet":
|
if task["status"] == "goal-unmet":
|
||||||
@@ -2423,6 +2468,13 @@ def main():
|
|||||||
if args.cmd == "dispatch":
|
if args.cmd == "dispatch":
|
||||||
if budget_exceeded(state, cfg):
|
if budget_exceeded(state, cfg):
|
||||||
print("预算已用完(max_daily_turns),本次不派发,见 doctor 的 usage")
|
print("预算已用完(max_daily_turns),本次不派发,见 doctor 的 usage")
|
||||||
|
# 前置预算门以前只打印:配了 notify_command 的
|
||||||
|
# 用户收不到任何通知。补一次 budget-exceeded
|
||||||
|
# 通知;没配命令时 notify_failure 只写日志,
|
||||||
|
# 不会无配置刷屏。
|
||||||
|
notify_failure(args.state, cfg, args.name,
|
||||||
|
{"status": "budget-exceeded",
|
||||||
|
"text": "budget exceeded: max_daily_turns"})
|
||||||
return
|
return
|
||||||
deps = [d.strip() for d in args.depends_on.split(",")
|
deps = [d.strip() for d in args.depends_on.split(",")
|
||||||
if d.strip()] if args.depends_on else []
|
if d.strip()] if args.depends_on else []
|
||||||
|
|||||||
@@ -262,6 +262,72 @@ class RetryAndBudgetTest(unittest.TestCase):
|
|||||||
os.path.join(tmp, "host-run.log")))
|
os.path.join(tmp, "host-run.log")))
|
||||||
self.assertIn("t1 failed", Path(notify_log).read_text())
|
self.assertIn("t1 failed", Path(notify_log).read_text())
|
||||||
|
|
||||||
|
def test_notify_failure_quotes_shell_metacharacters(self):
|
||||||
|
# 任务名 / 文本里的 shell 元字符必须被 quote,不能当命令执行:
|
||||||
|
# 未加 quote 时 task 名里的 `; touch pwn` 会真的建出文件。
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
state_path = os.path.join(tmp, "host-state.json")
|
||||||
|
notify_log = os.path.join(tmp, "notify.log")
|
||||||
|
pwn = os.path.join(tmp, "pwn")
|
||||||
|
pwn2 = os.path.join(tmp, "pwn2")
|
||||||
|
cfg = {"notify_command":
|
||||||
|
f"echo {{task}} {{status}} {{text}} >> {notify_log}"}
|
||||||
|
host.notify_failure(
|
||||||
|
state_path, cfg, "t1; touch " + pwn,
|
||||||
|
{"status": "failed", "text": "$(touch " + pwn2 + ")"})
|
||||||
|
self.assertFalse(os.path.exists(pwn))
|
||||||
|
self.assertFalse(os.path.exists(pwn2))
|
||||||
|
self.assertIn("t1; touch", Path(notify_log).read_text())
|
||||||
|
|
||||||
|
def test_notify_failure_bad_template_is_logged_not_swallowed(self):
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
state_path = os.path.join(tmp, "host-state.json")
|
||||||
|
notify_log = os.path.join(tmp, "notify.log")
|
||||||
|
cfg = {"notify_command": f"echo {{unknown}} >> {notify_log}"}
|
||||||
|
host.notify_failure(state_path, cfg, "t1",
|
||||||
|
{"status": "failed", "text": "boom"})
|
||||||
|
# 模板错了:不跑命令(notify.log 不应出现),但必须
|
||||||
|
# 在 host-run.log 里留下 template error,而不是静默吞掉。
|
||||||
|
self.assertFalse(os.path.exists(notify_log))
|
||||||
|
self.assertIn("notify command template error",
|
||||||
|
Path(os.path.join(tmp, "host-run.log")).read_text())
|
||||||
|
|
||||||
|
def test_notify_failure_command_failure_is_logged(self):
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
state_path = os.path.join(tmp, "host-state.json")
|
||||||
|
cfg = {"notify_command": "exit 3"}
|
||||||
|
host.notify_failure(state_path, cfg, "t1",
|
||||||
|
{"status": "failed", "text": "boom"})
|
||||||
|
self.assertIn("notify command rc=3",
|
||||||
|
Path(os.path.join(tmp, "host-run.log")).read_text())
|
||||||
|
|
||||||
|
def test_goal_budget_stop_notifies_budget_exceeded(self):
|
||||||
|
class FakeServer:
|
||||||
|
def run_turn(self, *args, **kwargs):
|
||||||
|
raise AssertionError("预算已满,不应再跑回合")
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
state_path = os.path.join(tmp, "host-state.json")
|
||||||
|
notify_log = os.path.join(tmp, "notify.log")
|
||||||
|
today = __import__("time").strftime("%Y-%m-%d")
|
||||||
|
state = {"tasks": {}, "automations": [],
|
||||||
|
"usage": {"date": today, "turns": 5}}
|
||||||
|
task = {"name": "g", "thread_id": "x", "cwd": tmp,
|
||||||
|
"sandbox": None, "max_retries": 0, "retry_count": 0,
|
||||||
|
"turns": []}
|
||||||
|
state["tasks"]["g"] = task
|
||||||
|
cfg = {"max_daily_turns": 5,
|
||||||
|
"notify_command":
|
||||||
|
f"echo {{task}} {{status}} >> {notify_log}"}
|
||||||
|
result = host.run_goal_loop(FakeServer(), state, state_path,
|
||||||
|
task, "make it", "true", 3, "m",
|
||||||
|
"low", None, cfg=cfg, timeout=600)
|
||||||
|
self.assertFalse(result["ok"])
|
||||||
|
self.assertTrue(result["budget_stopped"])
|
||||||
|
self.assertEqual(task["status"], "interrupted")
|
||||||
|
self.assertIn("budget-exceeded",
|
||||||
|
Path(notify_log).read_text())
|
||||||
|
|
||||||
|
|
||||||
class TurnErrorFromHistoryTest(unittest.TestCase):
|
class TurnErrorFromHistoryTest(unittest.TestCase):
|
||||||
def _server_with_history(self, tmp, error_json):
|
def _server_with_history(self, tmp, error_json):
|
||||||
|
|||||||
Reference in new issue
Block a user