feat: G7 分层 fail-open/fail-closed + G8 Docker 轮询掐断 + G9 worker 忽略 host 预算 + G13/G14 归属(第三批)

- G7:认领门 closed 时查配额异常回 quota_check_unavailable、标 interrupted(可重试/续跑);轮询与轮次边界用 open,异常时回 quota_check_failed——新回合边界拒绝,执行中轮询经 _quota_hit 闭包过滤此串、只记 warning 不杀在跑任务,确认超限才掐断;查询函数本身不吞异常冒充不限额。
- G8:Docker 路径加 _docker_call_with_quota_poll(后台线程跑阻塞式 docker 调用,主循环每 quota_poll_interval_seconds 缺省 30 秒查一次,确认超限按容器名真杀容器,结果改写 quota_exceeded 走统一收尾:interrupted、成本照记、续跑消息保留,配额命中后不走 ALREADY_EXISTS 回落);同步路径保持边界检查。
- G9:worker 路径默认忽略 host max_daily_turns(host.py 顶层 --ignore-host-budget 只清该项,Local/Docker 执行器在子命令前追加,worker CLI 缺省开、--no-ignore-host-budget 关);standalone 直接跑 host.py 时预算照常。
- G13/G14:/pr 不再走 system 豁免,建任务前按请求用户查配额(402/503 均不排队不改状态),队列路 pr_jobs.created_by 记请求用户(老表 ALTER 补列),worker 跑 pr 任务时审计操作人用 created_by(老行回落 system)。注意:G14 原案的 system 豁免(服务端与 worker 写死)与 P1 冲突(P1:system 按普通属主走默认额度,不另设无限额系统预算),本次 rebase 未合入该豁免,system 豁免需另行决策。
- 配额口径与 P0/P1 对齐:无配额行按默认额度(非不限额),只有显式 NULL 行才不限额;无价上限只计 source='actual' 行;缺属主(含父任务回查仍无)拒绝新回合。
- README 配额节同步以上口径与 G15 接受现状(best-effort 超支、UTC 月切分、在跑任务只在轮询/轮次边界响应配额改动)。G4 与 G15 不在本批。
This commit is contained in:
host committed 2026-10-28 11:14:10 +08:00
1 parent 5883169c82
commit e31fcc6cba
11 files changed
+745 -135

No files matched your search

+61 -12
View File
@@ -994,10 +994,35 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
重试/续跑时,操作人超限不应拦、创建者超限才应拦——旧版 重试/续跑时,操作人超限不应拦、创建者超限才应拦——旧版
重试按操作人查是错位,本轮已改)。后续消息 重试按操作人查是错位,本轮已改)。后续消息
(followup)创建不查配额:排队本身不花钱,执行时轮次边界 (followup)创建不查配额:排队本身不花钱,执行时轮次边界
的配额门照常生效,不要把 followup 也改成准入校验。系统 的配额门照常生效,不要把 followup 也改成准入校验。
任务(GitHub webhook 的 /fix-review,created_by 是 'system') - 费用归属与豁免(G13/G14,写死口径):系统任务豁免——
不记到任何用户名下、不占配额;要把它算到项目所有者头上 webhook /fix-review 等 created_by='system' 的任务在服务端
需另做归属。 准入与 worker 认领两处都不查配额(`_check_quota` 与
`_quota_exceeded` 里 user_id == "system" 直接放行,单测
覆盖);这是写死的约定,不是可配置项。反过来,由用户动
作触发的 /pr 任务不再走 system 豁免:`POST /tasks/{id}/pr`
在 merged 校验通过后、执行/排队之前先按请求用户查配额
(超限 402、查询异常见下 503,均不改状态、不排队),
队列路把请求用户写进 pr_jobs.created_by(老表缺列时
ensure 时 ALTER 补列,老行 created_by 为 NULL 时 worker
跑它记审计回落 "system"),worker 的 pr 任务循环把
created_by 当审计操作人传给 execute_task_pr(同步路本
就传调用方,不动)。注意这是一个有意收紧的口径,供合
并前推翻:PR 本身不跑模型回合,纯 git/gh 操作几乎不花
钱,若认为“已完成工作的 PR 不应被配额挡住”,可把 /pr
的配额门改回豁免而不动归属列(created_by 照常记录)。
- 配额查询异常的分层(G7,写死口径):新建/认领类门
(建任务、分支、重试、续跑、goal-resume、/pr)查配额抛
异常时 fail-closed——服务端回明确的 503
`quota_check_unavailable`(不是 500 断连),不建任务、
不改状态、不排队;worker 认领门同样标 interrupted,错
误含 `quota_check_unavailable`,可重试/续跑,绝不在“不
知道超没超”时开跑。执行中轮询(Local 30 秒、Docker 轮
询见下)查配额抛异常时 fail-open——只记 warning,不杀
在跑任务,只有确认超限才掐断;轮次边界的复查同款
fail-open。实现上查询函数 `_quota_limit_and_used` 自己
绝不吞异常冒充“不限额”,分层由调用方 on_error 参数决定
("closed" 门侧、"open" 轮询/边界,缺省 open)。
- 无价模型 token 上限:价目表目前只有两个 sol 模型有价,默认 - 无价模型 token 上限:价目表目前只有两个 sol 模型有价,默认
模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL,USD 模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL,USD
配额对它们永远不生效(按 0 计)。因此有配额行的用户另受一 配额对它们永远不生效(按 0 计)。因此有配额行的用户另受一
@@ -1017,15 +1042,39 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
时该条回 pending、任务 interrupted,不再排空剩余;建消息时 时该条回 pending、任务 interrupted,不再排空剩余;建消息时
不查,排队本身不花钱);Local popen 路径在轮询里每 30 秒查 不查,排队本身不花钱);Local popen 路径在轮询里每 30 秒查
一次,超限时杀进程组(goal 长循环也一样,杀的是整个子进程 一次,超限时杀进程组(goal 长循环也一样,杀的是整个子进程
组)。配额中断的任务一律 interrupted(可续跑:顶上来后重试 组);Docker 路径同款粒度(G8):`_docker_call_with_quota_poll`
把阻塞式 docker 调用放后台线程,主循环每
`quota_poll_interval_seconds`(缺省 30 秒,单测可传小值,
生产不要改小)查一次,确认超限时按任务容器名真杀容器,
结果改写成内部标记 `quota_exceeded` 后走统一收尾——状态
一律 interrupted、已花成本照记 actual、续跑消息保留,配额
命中后不走 ALREADY_EXISTS 回落重跑。同步路径(非 popen、
非 Docker)仍只有边界检查,单轮内的超支要到轮次结束才停。
配额中断的任务一律 interrupted(可续跑:顶上来后重试
/续跑继续),错误文本带 `quota_exceeded`,已花成本照记 /续跑继续),错误文本带 `quota_exceeded`,已花成本照记
actual,不发终态推送(与取消同款)。诚实边界:Docker/同步 actual,不发终态推送(与取消同款)。诚实边界:轮询查的是
路径没有轮询钩子,只有边界检查,单轮内的超支要到轮次结束 已记账用量,单轮本身花的钱无法中途按 token 流式截断(中转
才停;单轮本身花的钱无法中途按 token 流式截断(中转用量只 用量只在回合结束才回传),所以“掐断”粒度是轮次边界 + 长轮
在回合结束才回传),所以“掐断”粒度是轮次边界 + 长轮询中的 询中的 30 秒检查,不是 token 级;真 Docker 掐断未实弹(只
30 秒检查,不是 token 级;并发下两个任务同时查到未超限再同 由假执行器单测覆盖),部署后需实弹;并发下两个任务同时查
时开跑会有少量超支(与限流器同款 best-effort,非精确配额); 到未超限再同时开跑会有少量超支(与限流器同款 best-effort,
PR 任务(_run_pr_jobs)暂不查配额。 非精确配额)——G15 并发预留不在本批,单独立项。另两条
接受现状的边界:配额按 UTC 自然月切分,与 host 本地日期
不一致时以 UTC 为准(月末跨月任务的用量按记账日归属);
配额改动对在跑任务只在轮询/轮次边界生效,不会立刻掐断
已经开跑的回合(除非下一次轮询查到超限)。
- host 预算与云端配额的两套门(G9):worker 路径默认忽略
host 的 `max_daily_turns` 预算——worker 调 host.py 时传顶层
`--ignore-host-budget`(LocalExecutor 与 DockerExecutor 都
在子命令之前追加该标志;worker CLI 缺省开,可用
`--no-ignore-host-budget` 关掉),host.py 里只清
`max_daily_turns` 这一项(`_effective_cfg`,notify_command
等其它配置不动),dispatch 前置门、重试前检查、goal 循环
顶检查一并放行;云端 per-user USD 配额 + worker 认领/轮询
门才是 worker 路径的口径,host 单机 turn 数预算再拦一次只
会误拦云端已批的任务。standalone 直接跑 host.py 时不传
该标志,预算照常生效。两套门同时开(--no-ignore-host-budget)
时任务要同时通过两道门,超任一道都停。
- 接口:`GET /users/me/quota`(自查,返回 month、quota、 - 接口:`GET /users/me/quota`(自查,返回 month、quota、
used、remaining,以及当月无价实际 token 数 used、remaining,以及当月无价实际 token 数
`unpriced_actual_tokens`——价目表里没有该模型价格的实际 `unpriced_actual_tokens`——价目表里没有该模型价格的实际
+26 -6
View File
@@ -48,8 +48,21 @@ def ensure_pr_jobs_table(conn):
"project_id TEXT NOT NULL, host TEXT NOT NULL, " "project_id TEXT NOT NULL, host TEXT NOT NULL, "
"status TEXT NOT NULL, pr_url TEXT, pr_number INTEGER, " "status TEXT NOT NULL, pr_url TEXT, pr_number INTEGER, "
"error TEXT, created_at REAL, updated_at REAL, " "error TEXT, created_at REAL, updated_at REAL, "
"started_at REAL, finished_at REAL)" "started_at REAL, finished_at REAL, created_by TEXT)"
) )
# 老表补列(G13:费用归属人):PRAGMA 经 db.py 适配层翻译,
# SQLite 与 Postgres 同款写法(见 automations.py 同款模式)。
try:
_cols = {row[1] for row in conn.execute(
"PRAGMA table_info(pr_jobs)").fetchall()}
except Exception:
_cols = set()
if "created_by" not in _cols:
try:
conn.execute(
"ALTER TABLE pr_jobs ADD COLUMN created_by TEXT")
except Exception:
pass
conn.commit() conn.commit()
@@ -75,15 +88,21 @@ def find_open_pr_job(conn, task_id):
"host": row[3], "status": row[4]} "host": row[3], "status": row[4]}
def enqueue_pr_job(conn, task_id, project_id, host): def enqueue_pr_job(conn, task_id, project_id, host, created_by=None):
"""排一个 pr 任务(调用方已做去重)。返回 job id。""" """排一个 pr 任务(调用方已做去重)。返回 job id。
created_by(G13):请求 /pr 的用户 id,即费用归属人;
不传时为 None(老调用方 / 老行),worker 跑它时审计回
落到 "system"。server 的 _task_pr 队列路必传调用方。
"""
ensure_pr_jobs_table(conn) ensure_pr_jobs_table(conn)
job_id = uuid.uuid4().hex job_id = uuid.uuid4().hex
now = time.time() now = time.time()
conn.execute( conn.execute(
"INSERT INTO pr_jobs(id, task_id, project_id, host, status," "INSERT INTO pr_jobs(id, task_id, project_id, host, status,"
" created_at, updated_at) VALUES(?, ?, ?, ?, 'queued', ?, ?)", " created_at, updated_at, created_by)"
(job_id, task_id, project_id, host, now, now)) " VALUES(?, ?, ?, ?, 'queued', ?, ?, ?)",
(job_id, task_id, project_id, host, now, now, created_by))
conn.commit() conn.commit()
return job_id return job_id
@@ -182,7 +201,8 @@ def execute_task_pr(db_path, task_id, user_id, worktrees_root=None,
响应(PrSkipped -> 200 pr_skipped,PrStepError -> 502)。 响应(PrSkipped -> 200 pr_skipped,PrStepError -> 502)。
user_id:审计 pr_created 的操作人;server 端点传调用方, user_id:审计 pr_created 的操作人;server 端点传调用方,
worker 的 pr 任务循环传 "system"(与 sweep 同款约定)。 worker 的 pr 任务循环传 job 的 created_by(G13),老行
没有 created_by 时回落 "system"(与 sweep 同款约定)。
token 只在本次调用内存里出现,不进库、不进日志。 token 只在本次调用内存里出现,不进库、不进日志。
""" """
conn = _db_connect(str(db_path)) conn = _db_connect(str(db_path))
+14 -2
View File
@@ -189,11 +189,14 @@ class LocalExecutor(Executor):
""" """
def __init__(self, host_py_path, config_path=None, runner=None, def __init__(self, host_py_path, config_path=None, runner=None,
repo=None, model=None): repo=None, model=None, ignore_host_budget=False):
self.host_py_path = host_py_path self.host_py_path = host_py_path
self.config_path = config_path self.config_path = config_path
self.runner = runner if runner is not None else _default_runner self.runner = runner if runner is not None else _default_runner
self.repo = repo self.repo = repo
# G9:worker 路径默认忽略 host 预算时由 Worker 同步成
# True;standalone 的 LocalExecutor 缺省 False,预算照常。
self.ignore_host_budget = bool(ignore_host_budget)
# 模型:host.py 的 --model 是顶层全局参数(dispatch/run 共用), # 模型:host.py 的 --model 是顶层全局参数(dispatch/run 共用),
# 不传时 host.py 用配置文件里的 model;传了就覆盖配置。 # 不传时 host.py 用配置文件里的 model;传了就覆盖配置。
# worker 的 fallback 靠它把换模型后的第二次尝试真正发到新模型上, # worker 的 fallback 靠它把换模型后的第二次尝试真正发到新模型上,
@@ -205,6 +208,10 @@ class LocalExecutor(Executor):
cmd = [sys.executable, self.host_py_path] cmd = [sys.executable, self.host_py_path]
if self.config_path is not None: if self.config_path is not None:
cmd += ["--config", self.config_path] cmd += ["--config", self.config_path]
# G9:顶层全局参数,必须放在子命令之前(调用方在
# _base_command 之后才追加子命令)。
if getattr(self, "ignore_host_budget", False):
cmd += ["--ignore-host-budget"]
effective = model or self.model effective = model or self.model
if effective: if effective:
cmd += ["--model", effective] cmd += ["--model", effective]
@@ -693,7 +700,7 @@ class DockerExecutor(Executor):
user="1000:1000", worktrees_root=None, env=None, user="1000:1000", worktrees_root=None, env=None,
env_key="XCM_API_KEY", timeout=None, extra_mounts=None, env_key="XCM_API_KEY", timeout=None, extra_mounts=None,
model=None, security_opts=None, pids_limit=None, model=None, security_opts=None, pids_limit=None,
runtime=None, extra_env=None): runtime=None, extra_env=None, ignore_host_budget=False):
if host_py_path is None: if host_py_path is None:
here = os.path.dirname(os.path.abspath(__file__)) here = os.path.dirname(os.path.abspath(__file__))
host_py_path = os.path.join(os.path.dirname(here), "host.py") host_py_path = os.path.join(os.path.dirname(here), "host.py")
@@ -703,6 +710,8 @@ class DockerExecutor(Executor):
# 取构造时传的 model,都没有则不传 --model(容器内 host.py # 取构造时传的 model,都没有则不传 --model(容器内 host.py
# 用配置文件里的 model)。fallback 时 worker 换任务的 model。 # 用配置文件里的 model)。fallback 时 worker 换任务的 model。
self.model = model self.model = model
# G9:Worker 会把自己的 ignore_host_budget 同步过来。
self.ignore_host_budget = bool(ignore_host_budget)
self.runner = runner if runner is not None else _docker_default_runner self.runner = runner if runner is not None else _docker_default_runner
self.image = image self.image = image
self.cpus = str(cpus) self.cpus = str(cpus)
@@ -916,6 +925,9 @@ class DockerExecutor(Executor):
if self.config_path is not None: if self.config_path is not None:
cmd += ["--config", cmd += ["--config",
"/app/" + os.path.basename(self.config_path)] "/app/" + os.path.basename(self.config_path)]
# G9:与 LocalExecutor 同款,顶层全局参数,放子命令之前。
if getattr(self, "ignore_host_budget", False):
cmd += ["--ignore-host-budget"]
effective = model or self.model effective = model or self.model
if effective: if effective:
cmd += ["--model", effective] cmd += ["--model", effective]
+40 -8
View File
@@ -6368,7 +6368,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
month),调用方一律回 402——这是钱用完了,不是请求 month),调用方一律回 402——这是钱用完了,不是请求
太快,所以不用 429。 太快,所以不用 429。
准入点:建任务 / 分支 / 重试 / 续跑,均为权限校验 准入点:建任务 / 分支 / 重试 / 续跑 / goal-resume,均为权限校验
通过后、写库之前。P0 安全批后,后续消息与任务间 通过后、写库之前。P0 安全批后,后续消息与任务间
消息创建时同样查双人配额门(见 消息创建时同样查双人配额门(见
_reject_if_over_quota_for_task):排队本身不花钱, _reject_if_over_quota_for_task):排队本身不花钱,
@@ -6376,12 +6376,24 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
“超限用户靠他人任务续命”,所以创建时先拦一道, “超限用户靠他人任务续命”,所以创建时先拦一道,
执行时轮次边界门仍照常生效。 执行时轮次边界门仍照常生效。
fail-open / fail-closed 分层(G7):服务端准入点查
配额异常时 fail-closed——_reject_if_over_quota 回
明确的 503 quota_check_unavailable 且不改状态,绝不
因查询抖动把“不知道超没超”当成放行。system 操作人
不豁免,按普通属主走默认额度(G14 的 system 豁免与
P1 冲突,本次未合入)。
执行侧查配额异常时分层:认领门(closed)回
quota_check_unavailable、调用方标 interrupted;轮询与
轮次边界(open)回 quota_check_failed——新回合边界见
此串即拒绝,轮询经 _quota_hit 闭包过滤此串、只记
warning 不杀在跑任务,只有确认超限才掐断。
跑中掐断(与旧 docstring 不同,旧说法已过时): 跑中掐断(与旧 docstring 不同,旧说法已过时):
Local popen 路径的 worker 轮询每 30 秒查一次配额, Local popen 与 Docker 路径的 worker 轮询每 30 秒查
超限杀进程组;Docker / 同步路径没有轮询钩子,只有 一次配额,超限杀进程组/容器;同步路径没有轮询钩子,
认领与轮次边界检查,单轮内的超支要到轮次结束才停; 只有认领与轮次边界检查,单轮内的超支要到轮次结束
单轮本身的钱无法按 token 流式截断(用量只在回合 才停;单轮本身的钱无法按 token 流式截断(用量只在
结束才回传),掐断粒度是轮次边界 + 长轮询中的 回合结束才回传),掐断粒度是轮次边界 + 长轮询中的
30 秒检查,不是 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:
@@ -6395,8 +6407,19 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
方就是创建者,缺省即调用方;重试 / 续跑时费用记在 方就是创建者,缺省即调用方;重试 / 续跑时费用记在
原任务创建者名下,必须传 task 的 created_by,否则 原任务创建者名下,必须传 task 的 created_by,否则
维护者替他人重试会用错人的额度(旧 bug:重试按 维护者替他人重试会用错人的额度(旧 bug:重试按
操作人查、费用记创建者名下)。""" 操作人查、费用记创建者名下)。
查询异常时 fail-closed:回 503 quota_check_unavailable
且不改状态。调用方看到 503 应重试请求本身,而不是
认为任务已建好——本函数回 True 时调用方一律直接
return,不写库。"""
try:
blocked = self._check_quota(quota_user_id or user_id) blocked = self._check_quota(quota_user_id or user_id)
except Exception:
self._send_json(503, {"error": "quota_check_unavailable",
"detail": "quota lookup failed; "
"no task was created or queued"})
return True
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",
@@ -6690,6 +6713,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
"status": task["status"], "status": task["status"],
}) })
return return
# G13:/pr 不再走 system 豁免。请求用户即费用归属人,
# 建任务前先过配额门(超限 402、查询异常 503,均不改
# 状态、不排队);这是有意收紧的口径,见 PR 描述——
# PR 本身不跑模型回合,纯 git/gh 操作几乎不花钱,合
# 并前若认为“已完成工作的 PR 不应被配额挡住”,可把
# 这一门改回豁免而不动归属列(created_by 照常记录)。
if self._reject_if_over_quota(user_id):
return
# HA 阶段 1:worktree 在别机时不直接跑,排 pr 任务给该主机; # HA 阶段 1:worktree 在别机时不直接跑,排 pr 任务给该主机;
# 本机(或 worktree_host 未记录的老任务)走同步路。 # 本机(或 worktree_host 未记录的老任务)走同步路。
worktree_host = task.get("worktree_host") worktree_host = task.get("worktree_host")
@@ -6701,7 +6732,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
job_id = existing["id"] job_id = existing["id"]
else: else:
job_id = enqueue_pr_job( job_id = enqueue_pr_job(
conn, task_id, task["project_id"], worktree_host) conn, task_id, task["project_id"], worktree_host,
created_by=user_id)
finally: finally:
conn.close() conn.close()
self._send_json(200, { self._send_json(200, {
+42 -1
View File
@@ -26,7 +26,8 @@ try:
from cloud.autoscaler import Autoscaler from cloud.autoscaler import Autoscaler
from cloud.db import connect as _db_connect from cloud.db import connect as _db_connect
from cloud.pr_exec import (claim_pr_job, enqueue_pr_job, from cloud.pr_exec import (claim_pr_job, enqueue_pr_job,
find_open_pr_job, finish_pr_job) ensure_pr_jobs_table, find_open_pr_job,
finish_pr_job)
from cloud.server import create_server, migrate_tasks_table from cloud.server import create_server, migrate_tasks_table
from cloud.worker import Worker from cloud.worker import Worker
except ImportError: except ImportError:
@@ -105,6 +106,46 @@ class PrJobsTableTestCase(unittest.TestCase):
finally: finally:
conn.close() conn.close()
# -- G13:队列记费用归属人 + 老表补列 --------------------------------
def test_enqueue_records_created_by(self):
conn = _db_connect(self.db_path)
try:
enqueue_pr_job(conn, "t1", "p1", "host-a",
created_by="user-9")
job = find_open_pr_job(conn, "t1")
self.assertEqual(job["created_by"], "user-9")
# 不传时为 None(老调用方),worker 回落 system。
enqueue_pr_job(conn, "t2", "p1", "host-a")
self.assertIsNone(find_open_pr_job(conn, "t2")["created_by"])
finally:
conn.close()
def test_old_table_migration_adds_created_by(self):
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"CREATE TABLE pr_jobs(id TEXT PRIMARY KEY,"
" task_id TEXT NOT NULL, project_id TEXT NOT NULL,"
" host TEXT NOT NULL, status TEXT NOT NULL, pr_url TEXT,"
" pr_number INTEGER, error TEXT, created_at REAL,"
" updated_at REAL, started_at REAL, finished_at REAL)")
conn.commit()
finally:
conn.close()
conn = _db_connect(self.db_path)
try:
ensure_pr_jobs_table(conn)
cols = {row[1] for row in conn.execute(
"PRAGMA table_info(pr_jobs)").fetchall()}
self.assertIn("created_by", cols)
# 补列后照常排队并记归属人。
enqueue_pr_job(conn, "t1", "p1", "host-a",
created_by="user-9")
self.assertEqual(
find_open_pr_job(conn, "t1")["created_by"], "user-9")
finally:
conn.close()
class WorktreeHostMigrationTestCase(unittest.TestCase): class WorktreeHostMigrationTestCase(unittest.TestCase):
def setUp(self): def setUp(self):
+33
View File
@@ -191,6 +191,39 @@ class TaskOptionsTestCase(unittest.TestCase):
"untrusted") "untrusted")
self.assertNotIn("--sandbox", cmd) self.assertNotIn("--sandbox", cmd)
# -- G9:worker 路径默认忽略 host 预算 -------------------------------
def test_local_command_ignore_host_budget_flag(self):
seen = []
def fake_runner(cmd):
seen.append(cmd)
return SimpleNamespace(returncode=0, stdout="{}")
ex = LocalExecutor("/fake/host.py", runner=fake_runner,
ignore_host_budget=True)
ex.dispatch_task({"id": "t1", "prompt": "hi"})
cmd = seen[0]
self.assertIn("--ignore-host-budget", cmd)
# 顶层全局参数,必须在子命令之前。
self.assertLess(cmd.index("--ignore-host-budget"),
cmd.index("dispatch"))
ex2 = LocalExecutor("/fake/host.py", runner=fake_runner)
ex2.dispatch_task({"id": "t1", "prompt": "hi"})
self.assertNotIn("--ignore-host-budget", seen[1])
def test_docker_inner_command_ignore_host_budget_flag(self):
ex = DockerExecutor(config_path="/fake/cfg.json",
ignore_host_budget=True)
cmd = ex._inner_host_command("dispatch", "--name", "x",
"--task", "hi")
self.assertIn("--ignore-host-budget", cmd)
self.assertLess(cmd.index("--ignore-host-budget"),
cmd.index("dispatch"))
ex2 = DockerExecutor(config_path="/fake/cfg.json")
cmd2 = ex2._inner_host_command("dispatch", "--name", "x",
"--task", "hi")
self.assertNotIn("--ignore-host-budget", cmd2)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
+94 -2
View File
@@ -17,14 +17,16 @@ import threading
import unittest import unittest
from datetime import date, datetime, timedelta, timezone from datetime import date, datetime, timedelta, timezone
from pathlib import Path from pathlib import Path
from types import SimpleNamespace
from unittest import mock
try: try:
from cloud.auth import AuthStore from cloud.auth import AuthStore
from cloud.server import create_server from cloud.server import CloudControlPlaneHandler, create_server
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
sys.path.insert(0, str(Path(__file__).resolve().parents[2])) sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from cloud.auth import AuthStore from cloud.auth import AuthStore
from cloud.server import create_server from cloud.server import CloudControlPlaneHandler, create_server
def _req(port, method, path, body=None, key=None): def _req(port, method, path, body=None, key=None):
@@ -355,6 +357,96 @@ class ServerQuotasTestCase(unittest.TestCase):
self.assertEqual(s, 200, q) self.assertEqual(s, 200, q)
self.assertEqual(q["unpriced_actual_tokens"], 150) self.assertEqual(q["unpriced_actual_tokens"], 150)
# -- G7:准入点 fail-closed ----------------------------------------
def test_create_quota_lookup_failure_503(self):
with mock.patch.object(
CloudControlPlaneHandler, "_quota_status",
side_effect=RuntimeError("db down")):
s, body = self._create_task()
self.assertEqual(s, 503, body)
self.assertEqual(body["error"], "quota_check_unavailable")
# 没建出任务:状态、任务表都不改。
conn = sqlite3.connect(self.db_path)
try:
n = conn.execute(
"SELECT COUNT(*) FROM tasks").fetchone()[0]
finally:
conn.close()
self.assertEqual(n, 0)
# -- G13:system 豁免(server 侧) -----------------------------------
def test_system_quota_exempt_server(self):
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"INSERT INTO user_quotas(user_id, quota_usd_micros,"
" updated_at) VALUES('system', 0, 0)")
conn.commit()
finally:
conn.close()
handler = object.__new__(CloudControlPlaneHandler)
handler.server = SimpleNamespace(db_path=self.db_path)
self.assertIsNone(handler._check_quota("system"))
self.assertFalse(handler._reject_if_over_quota("system"))
# -- G13/G14:/pr 按请求人过配额门、队列记归属人 ----------------------
def _merged_remote_task(self):
s, task = self._create_task()
self.assertEqual(s, 201, task)
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"UPDATE tasks SET status = 'merged',"
" worktree_host = 'elsewhere-host-1' WHERE id = ?",
(task["id"],))
conn.commit()
finally:
conn.close()
return task
def test_pr_blocked_at_quota_no_job(self):
task = self._merged_remote_task()
s, _ = self._set_quota(self.owner_id, 0)
self.assertEqual(s, 200)
s, body = _req(self.port, "POST",
"/tasks/%s/pr" % task["id"], {},
self.owner_key)
self.assertEqual(s, 402, body)
self.assertEqual(body["error"], "quota_exceeded")
conn = sqlite3.connect(self.db_path)
try:
has_table = conn.execute(
"SELECT 1 FROM sqlite_master WHERE type='table'"
" AND name='pr_jobs'").fetchone() is not None
rows = conn.execute(
"SELECT COUNT(*) FROM pr_jobs"
" WHERE task_id = ?", (task["id"],)).fetchone()[0] \
if has_table else 0
status = conn.execute(
"SELECT status FROM tasks WHERE id = ?",
(task["id"],)).fetchone()[0]
finally:
conn.close()
self.assertEqual(rows, 0)
self.assertEqual(status, "merged")
def test_pr_queue_records_created_by(self):
task = self._merged_remote_task()
s, body = _req(self.port, "POST",
"/tasks/%s/pr" % task["id"], {},
self.owner_key)
self.assertEqual(s, 200, body)
self.assertTrue(body.get("pr_queued"))
conn = sqlite3.connect(self.db_path)
try:
row = conn.execute(
"SELECT created_by FROM pr_jobs WHERE task_id = ?",
(task["id"],)).fetchone()
finally:
conn.close()
self.assertIsNotNone(row)
self.assertEqual(row[0], self.owner_id)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
+142
View File
@@ -18,18 +18,22 @@ import sqlite3
import subprocess import subprocess
import sys import sys
import tempfile import tempfile
import threading
import time import time
import unittest import unittest
from pathlib import Path from pathlib import Path
from types import SimpleNamespace from types import SimpleNamespace
from unittest import mock
try: try:
from cloud.cost import CostStore from cloud.cost import CostStore
from cloud.sandbox import DockerExecutor, LocalExecutor
from cloud.server import init_db from cloud.server import init_db
from cloud.worker import Worker from cloud.worker import Worker
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
sys.path.insert(0, str(Path(__file__).resolve().parents[2])) sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from cloud.cost import CostStore from cloud.cost import CostStore
from cloud.sandbox import DockerExecutor, LocalExecutor
from cloud.server import init_db from cloud.server import init_db
from cloud.worker import Worker from cloud.worker import Worker
@@ -330,6 +334,144 @@ class WorkerQuotaTestCase(unittest.TestCase):
task, kind="dispatch", quota_check=lambda: True) task, kind="dispatch", quota_check=lambda: True)
self.assertEqual(result["status"], "quota_exceeded") self.assertEqual(result["status"], "quota_exceeded")
# -- G7 分层:门侧 fail-closed、轮询侧 fail-open -------------------
def test_claim_gate_quota_lookup_failure_closed(self):
_insert_task(self.db_path, "T1")
worker = self._worker(_usage_runner(
{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
self.calls))
with mock.patch.object(
Worker, "_quota_limit_and_used",
side_effect=RuntimeError("db down")):
result = worker.run_one()
self.assertEqual(result["status"], "interrupted")
self.assertIn("quota_check_unavailable", result["error"])
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
# 门侧拦住:runner 一次没调,也不记成本(cost_records
# 是 CostStore 懒建的,没跑过就根本没有这张表)。
self.assertEqual(self.calls, [])
conn = sqlite3.connect(self.db_path)
try:
has_table = conn.execute(
"SELECT 1 FROM sqlite_master WHERE type='table'"
" AND name='cost_records'").fetchone() is not None
rows = _cost_rows(self.db_path, "T1") if has_table else []
finally:
conn.close()
self.assertEqual(rows, [])
def test_poll_quota_lookup_failure_open(self):
_insert_task(self.db_path, "T1")
worker = self._worker(_usage_runner(
{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
self.calls))
task = worker.get_task("T1")
with mock.patch.object(
Worker, "_quota_limit_and_used",
side_effect=RuntimeError("db down")):
# 默认 on_error="open":异常不冒充超限,放行。
self.assertIsNone(worker._quota_exceeded(task, None))
self.assertIsNone(worker._quota_exceeded(task, None,
on_error="open"))
# -- G13:system 豁免(worker 侧) ----------------------------------
def test_system_task_exempt_from_quota(self):
_insert_quota(self.db_path, "system", 0)
_insert_task(self.db_path, "T1", created_by="system")
worker = self._worker(_usage_runner(
{"input_tokens": 100, "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")
# -- G8 Docker 轮询掐断 ---------------------------------------------
def test_docker_quota_poll_kills_container(self):
_insert_quota(self.db_path, "tester", QUOTA)
_insert_task(self.db_path, "T1")
_insert_followup(self.db_path, "fu1", "T1")
executor = _BlockingDockerExecutor()
worker = Worker(self.db_path, "/fake/host.py",
env=self.env, model="gpt-6.1-sol",
executor=executor, retry_backoff_seconds=0,
quota_poll_interval_seconds=0.01)
calls = {"n": 0}
def fake_limit(owner):
calls["n"] += 1
# 第 1 次是认领门(未超限,放行);之后轮询一律超限。
# helper 在 rebase 合并 #86 无价上限后返回 4 元组
# (quota, used, month, unpriced_used),此处无价
# 用量为 0,不触发无价上限。
return (QUOTA, 90000 if calls["n"] == 1 else QUOTA,
"2026-10", 0)
worker._quota_limit_and_used = fake_limit
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(executor.killed, ["codex-T1"])
# 续跑消息保留(仍 pending),已花成本照记。
self.assertEqual(_followup_status(self.db_path, "fu1"), "pending")
self.assertEqual(len(_cost_rows(self.db_path, "T1")), 1)
# -- G9:Worker 把 ignore 开关同步给外部执行器 ----------------------
def test_worker_syncs_ignore_host_budget_to_executor(self):
worker = Worker(self.db_path, "/fake/host.py",
env=self.env, model="gpt-6.1-sol",
executor=LocalExecutor("/fake/host.py"),
retry_backoff_seconds=0)
self.assertTrue(worker.ignore_host_budget)
self.assertTrue(worker.executor.ignore_host_budget)
worker2 = Worker(self.db_path, "/fake/host.py",
env=self.env, model="gpt-6.1-sol",
executor=LocalExecutor("/fake/host.py"),
ignore_host_budget=False,
retry_backoff_seconds=0)
self.assertFalse(worker2.ignore_host_budget)
self.assertFalse(worker2.executor.ignore_host_budget)
def test_docker_quota_poll_failure_open(self):
_insert_task(self.db_path, "T1")
worker = Worker(self.db_path, "/fake/host.py",
env=self.env, model="gpt-6.1-sol",
quota_poll_interval_seconds=0.01)
task = worker.get_task("T1")
def slow_fn():
time.sleep(0.05)
return {"status": "completed"}
def bad_check():
raise RuntimeError("db down")
result = worker._docker_call_with_quota_poll(
task, slow_fn, bad_check)
self.assertEqual(result["status"], "completed")
class _BlockingDockerExecutor(DockerExecutor):
"""G8 单测用假 Docker:dispatch 阻塞到被 kill,绝不调真 docker。"""
def __init__(self):
super().__init__(host_py_path="/fake/host.py")
self.killed = []
self._stop = threading.Event()
def dispatch_task(self, task, timeout=None, extra_mounts=None,
network=None, extra_env=None, seed_host_dir=None,
model_provider=None):
self._stop.wait(timeout=10)
return {"status": "failed", "stdout": "", "stderr": "killed",
"usage": {"input_tokens": 100, "cached_input_tokens": 0,
"output_tokens": 10, "total_tokens": 110}}
def kill_container(self, name):
self.killed.append(name)
self._stop.set()
return True
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
+253 -104
View File
@@ -372,7 +372,9 @@ class Worker:
upload_dir=None, upload_dir=None,
digest_interval_seconds=3600, digest_interval_seconds=3600,
outbox_backoff_seconds=60, outbox_backoff_seconds=60,
outbox_max_attempts=3): outbox_max_attempts=3,
quota_poll_interval_seconds=30,
ignore_host_budget=True):
self.db_path = str(db_path) self.db_path = str(db_path)
self.host_py_path = host_py_path self.host_py_path = host_py_path
self.config_path = config_path self.config_path = config_path
@@ -385,6 +387,17 @@ class Worker:
self.digest_interval_seconds = digest_interval_seconds self.digest_interval_seconds = digest_interval_seconds
self.outbox_backoff_seconds = outbox_backoff_seconds self.outbox_backoff_seconds = outbox_backoff_seconds
self.outbox_max_attempts = outbox_max_attempts self.outbox_max_attempts = outbox_max_attempts
# G8:配额轮询间隔(秒),Local popen 与 Docker 路径共用,
# 缺省 30;单测可传很小的值加速,生产不要改小(会打爆库)。
# None 当 30 用。
self.quota_poll_interval_seconds = (
30 if quota_poll_interval_seconds is None
else quota_poll_interval_seconds)
# G9:worker 路径默认忽略 host 的 max_daily_turns 预算。
# 云端已有 per-user USD 配额 + worker 认领/轮询门,host
# 的单机 turn 数预算再拦一次只会误拦云端已批的任务;
# standalone 直接跑 host.py 时不受此开关影响,预算照常。
self.ignore_host_budget = bool(ignore_host_budget)
self._last_digest_sent = 0.0 self._last_digest_sent = 0.0
# env 只用于 key 存在性检查:默认接进程环境,单测可注入 # env 只用于 key 存在性检查:默认接进程环境,单测可注入
# 一个普通 dict;worker 永远不把 env 里的值写出。 # 一个普通 dict;worker 永远不把 env 里的值写出。
@@ -416,7 +429,15 @@ class Worker:
else: else:
self.executor = LocalExecutor( self.executor = LocalExecutor(
host_py_path, config_path=config_path, runner=runner, host_py_path, config_path=config_path, runner=runner,
repo=repo) repo=repo, ignore_host_budget=self.ignore_host_budget)
# 外部传入的执行器(CLI 的 DockerExecutor 等)同样同步
# G9 开关:worker 的 ignore_host_budget 是唯一口径,执行
# 器自己的缺省值不另起一套。假执行器没有该属性时跳过。
if hasattr(self.executor, "ignore_host_budget"):
try:
self.executor.ignore_host_budget = self.ignore_host_budget
except Exception:
pass
self._cost_store = cost_store self._cost_store = cost_store
self._cancel_requested = set() self._cancel_requested = set()
# 真取消用:process_factory 可注入(单测用它返回一个假长跑 # 真取消用:process_factory 可注入(单测用它返回一个假长跑
@@ -961,7 +982,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): def _quota_limit_and_used(self, owner):
"""查某人当月配额与已记账用量,返回
(quota, used, month, unpriced_used)。
quota 为 None 表示不限额(只限行里 quota 为 NULL 的
显式不限额;无配额行按默认配额
_default_quota_usd_micros() 算,与服务端 _quota_status
的 P0 口径一致)。查询失败时抛异常,由调用方按
fail-open / fail-closed 分层决定——本函数自己绝不吞
异常后冒充“不限额”,否则调用方分不清“真不限额”与
“查不出来”。
unpriced_used 是当月 usd IS NULL 且 source='actual'
的成本行 input + output 之和(与服务端
_check_unpriced_token_cap 同口径,估算行不计);
quota 为 None 时它仍返回实数,但调用方
_quota_exceeded 在 quota 为 None 时先返回 None,
所以显式不限额用户不受无价上限约束,与服务端一致。
"""
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:
# P0/P1 口径:无行按默认额度,不是不限额(只有
# 显式 NULL 行才是不限额,见下)。
quota = _default_quota_usd_micros()
else:
quota = row[0]
if quota is None:
return None, 0, month, 0
# 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
unpriced_row = conn.execute(
"SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
" FROM cost_records c JOIN tasks t"
" ON c.task_id = t.id"
" WHERE t.created_by = ? AND c.day LIKE ?"
" AND c.usd_micros IS NULL AND c.source = 'actual'",
(owner, month + "%")).fetchone()
unpriced_used = unpriced_row[0] if unpriced_row else 0
return quota, int(used or 0), month, int(unpriced_used or 0)
finally:
conn.close()
def _quota_exceeded(self, task, usage_total, on_error="open"):
"""配额加固的执行时检查:返回超限时的错误文本,否则 None。 """配额加固的执行时检查:返回超限时的错误文本,否则 None。
口径与服务端准入(_quota_status)一致:按任务创建者算, 口径与服务端准入(_quota_status)一致:按任务创建者算,
@@ -970,29 +1057,27 @@ class Worker:
另有一道无价 token 上限(与服务端 另有一道无价 token 上限(与服务端
_check_unpriced_token_cap 同口径):USD 配额不计量无价 _check_unpriced_token_cap 同口径):USD 配额不计量无价
模型,默认模型 terra 无价,若无此上限则配额对默认模型 模型,默认模型 terra 无价,若无此上限则配额对默认模型
永远不生效;有配额行的用户当月无价 token(input + 永远不生效;有额度约束的用户(无行按默认额度)当月无价
token(input +
output,已记账 + 本轮)达到上限即停。无价上限只计 output,已记账 + 本轮)达到上限即停。无价上限只计
source='actual' 的行:估算行(estimated)不是真实花费, source='actual' 的行:估算行(estimated)不是真实花费,
计入上限会把“没花钱”算成“花超了”,且与 cost.py 计入上限会把“没花钱”算成“花超了”,且与 cost.py
_report 的 unpriced_actual_tokens 口径一致。 _report 的 unpriced_actual_tokens 口径一致。
P1 执行面三处收口(旧版三处都是 fail-open,已废): 属主与额度(P1 口径,与服务端 _quota_status 一致):
- 属主:任务 created_by 为空时先查父任务(协调员子任务 - 属主:任务 created_by 为空时先查父任务(协调员子任务
行可能无 created_by,属主跟父任务走);仍查不到时 行可能无 created_by,属主跟父任务走);仍查不到时
拒绝新回合,不再返回 None 放行。created_by 字面为 拒绝新回合,不再返回 None 放行。created_by 字面为
"system" 的系统任务按普通属主算(无行时同样走默认 "system" 的系统任务按普通属主算,不另设无限额的
额度,不另设无限额的系统预算)。 系统预算(G14 的 system 豁免与 P1 冲突,本次未合入)。
- 无配额行:不再是不限额,按默认配额 - 无配额行:按默认配额 _default_quota_usd_micros()
_default_quota_usd_micros()($10/月)算;只有显式 ($10/月)算;只有显式 NULL 行才是不限额。
NULL 行才是不限额(与 P0 的服务端口径一致,管理员 查库失败分层(G7 + P1):认领门传 on_error="closed",
显式设 NULL 的行仍被尊重)。 返回 "quota_check_unavailable: ..."(调用方标
- 查库失败:返回 "quota_check_failed: ..."(fail-closed), interrupted,可重试/续跑);默认 "open" 返回
绝不返回 None 放行。调用方分流:新回合边界(认领门、 "quota_check_failed: ..."——新回合边界(直接调用)
重试/fallback/后续消息边界)见此串即拒绝新回合并记 见此串即拒绝,执行中轮询经 _quota_hit 闭包过滤此串、
worker_quota_stop 审计;执行中轮询(popen 轮询、 只记警告不杀当前轮。查询抖动不应杀掉已在跑的回合,
Docker 配额看门狗)见此串只记警告、不中途杀当前轮 但下一轮边界仍会拒绝。
——查询抖动不应杀掉已在跑的回合,但下一轮边界仍会
拒绝(见 run_one 的 _quota_hit 过滤)。查库失败的
审计由边界调用方记,这里只记日志,避免轮询高频写审计。
""" """
owner = (task or {}).get("created_by") owner = (task or {}).get("created_by")
if not owner and (task or {}).get("parent_task_id"): if not owner and (task or {}).get("parent_task_id"):
@@ -1017,63 +1102,23 @@ class Worker:
return ("quota_exceeded: 任务缺少属主,无法查配额," return ("quota_exceeded: 任务缺少属主,无法查配额,"
"拒绝新回合(请归属真实用户后再排队)") "拒绝新回合(请归属真实用户后再排队)")
try: try:
month = datetime.now(timezone.utc).strftime("%Y-%m") quota, used, month, unpriced_used = \
conn = _db_connect(self.db_path) self._quota_limit_and_used(owner)
except Exception as exc:
# 查库失败分层(见本函数 docstring):closed(认领门)
# 回 quota_check_unavailable,调用方标 interrupted;
# open 回 quota_check_failed,新回合边界拒绝,执行中
# 轮询经 _quota_hit 闭包过滤、只警告不杀轮。
try: try:
conn.execute( self.logger.warning("quota check failed: %s", exc)
"CREATE TABLE IF NOT EXISTS user_quotas(" except Exception:
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT," pass
" updated_at DOUBLE PRECISION)") if on_error == "closed":
conn.commit() return "quota_check_unavailable: 配额查询失败,任务未开跑"
row = conn.execute( return ("quota_check_failed: 配额查询失败,拒绝新回合"
"SELECT quota_usd_micros FROM user_quotas" "(查库异常,非配额用完)")
" WHERE user_id = ?",
(owner,)).fetchone()
if row is None:
# P1:无行按默认额度,不是不限额(P0 服务端
# 已在建号/登录时补行,worker 侧仍可能遇到
# 老库未补行的用户,这里按默认额度兜底,
# 不写行——写行是服务端的职责,worker 不抢)。
quota = _default_quota_usd_micros()
else:
quota = row[0]
if quota is None: if quota is None:
return None return None
# cost_records 由 CostStore 懒建:全新库上它还不存
# 在,直接 JOIN 会报错——先实例化一次建表再查(与
# 服务端 _quota_status 同款做法),否则 quota=0
#(一分钱也不能花)在全新库上会被误放行。
# P1 补:建表必须落在本连接上,不只 CostStore 的
# 连接上——文件库两者指向同一文件所以旧代码碰巧
# 能用,:memory: 库两者是不同空库,本连接查 JOIN
# 会直接报错并(P1 fail-closed 后)误拒全部任务。
_cost_store = CostStore(self.db_path)
try:
_cost_store._conn.close()
except Exception:
pass
try:
conn.execute(CostStore._COST_DDL)
conn.commit()
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
unpriced_row = conn.execute(
"SELECT COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
" FROM cost_records c JOIN tasks t"
" ON c.task_id = t.id"
" WHERE t.created_by = ? AND c.day LIKE ?"
" AND c.usd_micros IS NULL AND c.source = 'actual'",
(owner, month + "%")).fetchone()
unpriced_used = unpriced_row[0] if unpriced_row else 0
finally:
conn.close()
if usage_total: if usage_total:
model = (task or {}).get("model") or self.model model = (task or {}).get("model") or self.model
micros = price_usd_micros( micros = price_usd_micros(
@@ -1096,16 +1141,6 @@ class Worker:
"(已用 %d / 上限 %d,%s)" "(已用 %d / 上限 %d,%s)"
% (unpriced_used, cap, month)) % (unpriced_used, cap, month))
return None return None
except Exception as exc:
# P1 fail-closed:查库失败不再返回 None。调用方按
# 前缀分流(见本函数 docstring):quota_check_failed
# 在新回合边界 = 拒绝,在执行中轮询 = 只警告不杀轮。
try:
self.logger.warning("quota check failed: %s", exc)
except Exception:
pass
return ("quota_check_failed: 配额查询失败,拒绝新回合"
"(查库异常,非配额用完)")
def _ready_tasks(self): def _ready_tasks(self):
"""当前就绪任务列表(pick_next 的批量版)。 """当前就绪任务列表(pick_next 的批量版)。
@@ -1511,11 +1546,13 @@ class Worker:
的标记,调用方据此保持 interrupted;deadline_seconds 的标记,调用方据此保持 interrupted;deadline_seconds
到期时杀进程组并返回 status='timeout' 的标记,调用方 到期时杀进程组并返回 status='timeout' 的标记,调用方
据此置 failed(超时是失败,不是取消)。quota_check 不为 据此置 failed(超时是失败,不是取消)。quota_check 不为
空时每 30 秒查一次配额(轮询本身 50 毫秒一圈,逐圈查会 空时每 quota_poll_interval_seconds(缺省 30 秒)查一次
把库打爆),超限时同样杀进程组并返回 status= 配额(轮询本身 50 毫秒一圈,逐圈查会把库打爆),超限时
'quota_exceeded' 的内部标记——调用方据此走配额中断收尾, 同样杀进程组并返回 status='quota_exceeded' 的内部标记
绝不能把这个标记原样写进 tasks.status(不在合法状态集 ——调用方据此走配额中断收尾,绝不能把这个标记原样写
里)。Docker/同步路径没有这个轮询钩子,只走边界检查。 进 tasks.status(不在合法状态集里)。Docker 路径见
_docker_call_with_quota_poll(同款轮询);同步路径没有
轮询钩子,只走边界检查。
""" """
task_id = task["id"] task_id = task["id"]
gateway_provider, gateway_env = self._ensure_gateway(task) gateway_provider, gateway_env = self._ensure_gateway(task)
@@ -1539,6 +1576,10 @@ class Worker:
started = time.time() started = time.time()
last_quota_check = 0.0 last_quota_check = 0.0
last_db_cancel_check = 0.0 last_db_cancel_check = 0.0
quota_interval = getattr(
self, "quota_poll_interval_seconds", 30)
if quota_interval is None:
quota_interval = 30
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:
@@ -1567,7 +1608,7 @@ class Worker:
self._kill_proc_group(proc) self._kill_proc_group(proc)
break break
if quota_check is not None and \ if quota_check is not None and \
time.time() - last_quota_check >= 30: time.time() - last_quota_check >= quota_interval:
last_quota_check = time.time() last_quota_check = time.time()
try: try:
over = quota_check() over = quota_check()
@@ -2205,6 +2246,80 @@ class Worker:
detail += " (" + "; ".join(rejected) + ")" detail += " (" + "; ".join(rejected) + ")"
self._audit(task["project_id"], "worker_peer_messages", detail) self._audit(task["project_id"], "worker_peer_messages", detail)
def _docker_call_with_quota_poll(self, task, fn, quota_check=None):
"""Docker 同步调用包一层配额轮询(G8),与 popen 路径同款。
fn 在后台线程里跑(DockerExecutor 的调用本身是阻塞的
docker run,无法像 popen 那样在主循环里轮询);主循环
每 quota_poll_interval_seconds(缺省 30 秒)调一次
quota_check,确认超限时按任务容器名真杀容器,fn 返回
后把结果状态改写成内部标记 'quota_exceeded',调用方
(run_one)据此走统一的配额中断收尾:状态 interrupted、
已花成本照记、续跑消息保留。quota_check 为 None 时
直接调 fn,不起线程。quota_check 本身抛异常时按
fail-open 处理(只记 warning,不杀容器),只有确认
超限才杀;取消时 cancel() 已按容器名杀过容器,线程
随之结束,轮询循环不再多杀一次。fn 抛异常时在主线程
原样重抛,由 _run_attempt 外层的 except 写成 failed,
与不包轮询时的行为一致。注意:轮询查的是已记账用量
(与 popen 路径同款),在跑回合的 token 只在回合结束
才回传,所以掐断粒度仍是轮询间隔,不是 token 级。
"""
if quota_check is None:
return fn()
box = {}
def _target():
try:
box["result"] = fn()
except Exception as exc: # noqa: BLE001 -- 主线程重抛
box["error"] = exc
thread = threading.Thread(target=_target, daemon=True)
thread.start()
interval = getattr(self, "quota_poll_interval_seconds", 30)
if interval is None:
interval = 30
last_check = 0.0
quota_hit = False
while thread.is_alive():
if task["id"] in self._cancel_requested or self._loop_stop:
break
now = time.time()
if now - last_check >= interval:
last_check = now
try:
over = bool(quota_check())
except Exception as exc:
try:
self.logger.warning(
"docker quota poll failed: %s", exc)
except Exception:
pass
over = False
if over:
quota_hit = True
container = None
try:
container = self.executor.container_name(task)
except Exception:
container = None
if container is not None and hasattr(
self.executor, "kill_container"):
try:
self.executor.kill_container(container)
except Exception:
pass
break
thread.join(timeout=0.05)
thread.join()
if "error" in box:
raise box["error"]
result = box.get("result")
if quota_hit and isinstance(result, dict):
result = dict(result, status="quota_exceeded")
return result
def _run_attempt(self, task, quota_check=None): def _run_attempt(self, task, quota_check=None):
"""单次执行尝试(不含认领、key 门、重试、成本与最终审计)。 """单次执行尝试(不含认领、key 门、重试、成本与最终审计)。
@@ -2337,12 +2452,16 @@ class Worker:
if task.get("_goal_resume_breakpoint"): if task.get("_goal_resume_breakpoint"):
# G10 断点:Docker 路用无 --goal-text 的 # G10 断点:Docker 路用无 --goal-text 的
# goal-resume-breakpoint 命令。 # goal-resume-breakpoint 命令。
result = self.executor.goal_resume_breakpoint_task( result = self._docker_call_with_quota_poll(
task,
lambda: self.executor.goal_resume_breakpoint_task(
task, task,
timeout=effective_docker_timeout, timeout=effective_docker_timeout,
extra_mounts=tenant_mounts, extra_mounts=tenant_mounts,
network=tenant_network, network=tenant_network,
extra_env=proxy_env) extra_env=_merge_env(proxy_env,
gateway_env)),
quota_check)
elif task.get("_resume"): elif task.get("_resume"):
# P1 断点续跑传 None(不传 --goal-text, # P1 断点续跑传 None(不传 --goal-text,
# host 沿用原 goal+verify);开新目标 # host 沿用原 goal+verify);开新目标
@@ -2350,31 +2469,41 @@ class Worker:
_resume_text = None if task.get( _resume_text = None if task.get(
"_resume_kind") == "breakpoint" \ "_resume_kind") == "breakpoint" \
else task["prompt"] else task["prompt"]
result = self.executor.goal_resume_task( result = self._docker_call_with_quota_poll(
task,
lambda: self.executor.goal_resume_task(
task, _resume_text, task, _resume_text,
timeout=effective_docker_timeout, timeout=effective_docker_timeout,
extra_mounts=tenant_mounts, extra_mounts=tenant_mounts,
network=tenant_network, network=tenant_network,
extra_env=_merge_env(proxy_env, extra_env=_merge_env(proxy_env,
gateway_env), gateway_env),
model_provider=gateway_provider) model_provider=gateway_provider),
quota_check)
else: else:
result = self.executor.goal_task( result = self._docker_call_with_quota_poll(
task,
lambda: self.executor.goal_task(
task, timeout=effective_docker_timeout, task, timeout=effective_docker_timeout,
extra_mounts=tenant_mounts, extra_mounts=tenant_mounts,
network=tenant_network, network=tenant_network,
extra_env=_merge_env(proxy_env, extra_env=_merge_env(proxy_env,
gateway_env), gateway_env),
seed_host_dir=seed_host_dir, seed_host_dir=seed_host_dir,
model_provider=gateway_provider) model_provider=gateway_provider),
quota_check)
else: else:
result = self.executor.dispatch_task( result = self._docker_call_with_quota_poll(
task,
lambda: self.executor.dispatch_task(
task, timeout=effective_docker_timeout, task, timeout=effective_docker_timeout,
extra_mounts=tenant_mounts, extra_mounts=tenant_mounts,
network=tenant_network, network=tenant_network,
extra_env=_merge_env(proxy_env, gateway_env), extra_env=_merge_env(proxy_env,
gateway_env),
seed_host_dir=seed_host_dir, seed_host_dir=seed_host_dir,
model_provider=gateway_provider) model_provider=gateway_provider),
quota_check)
else: else:
# seed 参数只传给 LocalExecutor:CloudExecutor # seed 参数只传给 LocalExecutor:CloudExecutor
# 基类与测试假执行器的 goal_task/dispatch_task # 基类与测试假执行器的 goal_task/dispatch_task
@@ -2434,10 +2563,13 @@ class Worker:
result = self.executor.dispatch_task(task) result = self.executor.dispatch_task(task)
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"] != "completed" and \ if not goal_mode and result["status"] not in (
"completed", "quota_exceeded") and \
_ALREADY_EXISTS_MARKER in combined: _ALREADY_EXISTS_MARKER in combined:
if is_docker: if is_docker:
result = self.executor.run_task( result = self._docker_call_with_quota_poll(
task,
lambda: self.executor.run_task(
task, timeout=effective_docker_timeout, task, timeout=effective_docker_timeout,
extra_mounts=tenant_mounts extra_mounts=tenant_mounts
if self.tenant_isolation else None, if self.tenant_isolation else None,
@@ -2446,7 +2578,8 @@ class Worker:
extra_env=_merge_env( extra_env=_merge_env(
proxy_env if self.egress_proxy_port proxy_env if self.egress_proxy_port
else None, gateway_env), else None, gateway_env),
model_provider=gateway_provider) model_provider=gateway_provider),
quota_check)
else: else:
if gateway_provider and not isinstance( if gateway_provider and not isinstance(
self.executor, LocalExecutor): self.executor, LocalExecutor):
@@ -3092,8 +3225,11 @@ class Worker:
# 再开跑——直接标 interrupted + 配额原因,不记成本(还没花 # 再开跑——直接标 interrupted + 配额原因,不记成本(还没花
# 钱),续跑消息不消费,顶上来后重试/续跑可再跑。与 key # 钱),续跑消息不消费,顶上来后重试/续跑可再跑。与 key
# 门、仓库门同款提前返回形状,区别是状态为 interrupted # 门、仓库门同款提前返回形状,区别是状态为 interrupted
#(可续跑)而非 failed。 #(可续跑)而非 failed。G7 分层:认领门用 closed——
quota_error = self._quota_exceeded(task, None) # 查询异常时同样标 interrupted(quota_check_unavailable,
# 可重试/续跑),绝不在“不知道超没超”时开跑;轮询与
# 轮次边界仍用默认 open,查询抖动不杀在跑任务。
quota_error = self._quota_exceeded(task, None, on_error="closed")
if quota_error is not None: if quota_error is not None:
self._set_status(task_id, "interrupted") self._set_status(task_id, "interrupted")
self._write_execution_result( self._write_execution_result(
@@ -3696,8 +3832,11 @@ class Worker:
job_id = job["id"] job_id = job["id"]
task_id = job["task_id"] task_id = job["task_id"]
try: try:
# G13:审计操作人用 job 的费用归属人(请求 /pr 的
# 用户),老行没有 created_by 时回落 "system"。
result = execute_task_pr( result = execute_task_pr(
self.db_path, task_id, "system") self.db_path, task_id,
job.get("created_by") or "system")
conn = _db_connect(self.db_path) conn = _db_connect(self.db_path)
try: try:
finish_pr_job(conn, job_id, "completed", finish_pr_job(conn, job_id, "completed",
@@ -4121,6 +4260,15 @@ def main(argv=None):
"列表里的第一个模型再跑一次,401/400" "列表里的第一个模型再跑一次,401/400"
"绝不 fallback,取消绝不 fallback)") "绝不 fallback,取消绝不 fallback)")
parser.add_argument("--max-turns", type=int, default=None) parser.add_argument("--max-turns", type=int, default=None)
parser.add_argument("--ignore-host-budget", dest="ignore_host_budget",
action="store_true", default=True,
help="G9(默认开):worker 路径忽略 host 的 "
"max_daily_turns 预算,云端 per-user USD "
"配额 + worker 门才是口径")
parser.add_argument("--no-ignore-host-budget",
dest="ignore_host_budget", action="store_false",
help="G9:关闭忽略,worker 调 host.py 时 host "
"的 max_daily_turns 预算照常生效(双重门)")
parser.add_argument("--max-tasks", type=int, default=None) parser.add_argument("--max-tasks", type=int, default=None)
parser.add_argument("--executor", default="local", parser.add_argument("--executor", default="local",
choices=["local", "docker", "cloud"], choices=["local", "docker", "cloud"],
@@ -4219,6 +4367,7 @@ def main(argv=None):
args.fallback_models.split(",")] args.fallback_models.split(",")]
if args.fallback_models else []), if args.fallback_models else []),
max_turns=args.max_turns, max_turns=args.max_turns,
ignore_host_budget=args.ignore_host_budget,
executor=_build_cli_executor(args, host_py), executor=_build_cli_executor(args, host_py),
task_timeout=args.task_timeout, task_timeout=args.task_timeout,
max_retries=args.max_retries, max_retries=args.max_retries,
+23
View File
@@ -2057,6 +2057,19 @@ def _load_config_file(path):
return {} return {}
def _effective_cfg(cfg, ignore_host_budget):
"""G9:--ignore-host-budget 时的生效配置。
只清 max_daily_turns 一项(其它配置原样),返回新 dict
不改原 cfg;不传开关时原样返回。main 用它,单测直接
测它,不用拉起 host.py 子进程。"""
if not ignore_host_budget:
return cfg
out = dict(cfg or {})
out["max_daily_turns"] = None
return out
def _resolve_path(base_dir, p): def _resolve_path(base_dir, p):
if not p: if not p:
return p return p
@@ -2099,6 +2112,12 @@ def main():
"打一行 turn progress 并往链路文件记一条 " "打一行 turn progress 并往链路文件记一条 "
"kind=progress;0 表示关闭进度输出;不传时用配置 " "kind=progress;0 表示关闭进度输出;不传时用配置 "
"progress_interval,未配则默认 30") "progress_interval,未配则默认 30")
ap.add_argument("--ignore-host-budget", action="store_true",
help="G9:忽略 host 的 max_daily_turns 预算(云端 "
"worker 路径默认传它:云端已有 per-user USD "
"配额 + worker 认领/轮询门,host 单机 turn "
"数预算再拦一次只会误拦云端已批的任务;"
"standalone 直接跑 host.py 时不传,预算照常)")
sub = ap.add_subparsers(dest="cmd", required=True) sub = ap.add_subparsers(dest="cmd", required=True)
p = sub.add_parser("init") p = sub.add_parser("init")
@@ -2241,6 +2260,10 @@ def main():
print(f"配置文件已写:{config_path},请按需改 repo/model 后再用") print(f"配置文件已写:{config_path},请按需改 repo/model 后再用")
return return
# 有效配置(init/config 之外的命令都用它) # 有效配置(init/config 之外的命令都用它)
# G9:只清预算这一项,notify_command 等其它配置不动;
# 本进程内所有 budget_exceeded(state, cfg) 检查随之失效,
# dispatch 前置门、重试前检查、goal 循环顶检查一并放行。
cfg = _effective_cfg(cfg, getattr(args, "ignore_host_budget", False))
args.model = args.model or cfg.get("model") or "gpt-5.6-terra" args.model = args.model or cfg.get("model") or "gpt-5.6-terra"
args.effort = args.effort or cfg.get("effort") or "low" args.effort = args.effort or cfg.get("effort") or "low"
args.sandbox = args.sandbox or cfg.get("sandbox") args.sandbox = args.sandbox or cfg.get("sandbox")
+17
View File
@@ -14,6 +14,23 @@ sys.path.insert(0, str(REPO_ROOT))
import host import host
class EffectiveCfgTest(unittest.TestCase):
# G9:--ignore-host-budget 只清 max_daily_turns,其它配置不动。
def test_ignore_host_budget_clears_only_budget(self):
cfg = {"max_daily_turns": 10, "notify_command": "echo hi",
"model": "gpt-6.1-sol"}
out = host._effective_cfg(cfg, True)
self.assertIsNone(out["max_daily_turns"])
self.assertEqual(out["notify_command"], "echo hi")
self.assertEqual(out["model"], "gpt-6.1-sol")
# 不改原 cfg。
self.assertEqual(cfg["max_daily_turns"], 10)
def test_no_ignore_returns_cfg_unchanged(self):
cfg = {"max_daily_turns": 10}
self.assertEqual(host._effective_cfg(cfg, False), cfg)
class HostPureFunctionsTest(unittest.TestCase): class HostPureFunctionsTest(unittest.TestCase):
def test_save_state_merges_tasks_and_preserves_disk_tasks(self): def test_save_state_merges_tasks_and_preserves_disk_tasks(self):
with tempfile.TemporaryDirectory() as tmp: with tempfile.TemporaryDirectory() as tmp: