From e31fcc6cba2728e19cd5d31f8612e712a7e4bd36 Mon Sep 17 00:00:00 2001 From: host Date: Tue, 27 Oct 2026 10:07:44 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20G7=20=E5=88=86=E5=B1=82=20fail-open/fai?= =?UTF-8?q?l-closed=20+=20G8=20Docker=20=E8=BD=AE=E8=AF=A2=E6=8E=90?= =?UTF-8?q?=E6=96=AD=20+=20G9=20worker=20=E5=BF=BD=E7=95=A5=20host=20?= =?UTF-8?q?=E9=A2=84=E7=AE=97=20+=20G13/G14=20=E5=BD=92=E5=B1=9E=EF=BC=88?= =?UTF-8?q?=E7=AC=AC=E4=B8=89=E6=89=B9=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 不在本批。 --- cloud/README.md | 73 +++++- cloud/pr_exec.py | 32 ++- cloud/sandbox.py | 16 +- cloud/server.py | 50 +++- cloud/tests/test_pr_jobs.py | 43 ++- cloud/tests/test_task_options.py | 33 +++ cloud/tests/test_task_quotas.py | 96 ++++++- cloud/tests/test_worker_quota.py | 142 ++++++++++ cloud/worker.py | 437 +++++++++++++++++++++---------- host.py | 23 ++ tests/test_host.py | 17 ++ 11 files changed, 786 insertions(+), 176 deletions(-) diff --git a/cloud/README.md b/cloud/README.md index f727683..baed554 100644 --- a/cloud/README.md +++ b/cloud/README.md @@ -994,10 +994,35 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳 重试/续跑时，操作人超限不应拦、创建者超限才应拦——旧版 重试按操作人查是错位，本轮已改）。后续消息 （followup）创建不查配额：排队本身不花钱，执行时轮次边界 - 的配额门照常生效，不要把 followup 也改成准入校验。系统 - 任务（GitHub webhook 的 /fix-review，created_by 是 'system'） - 不记到任何用户名下、不占配额；要把它算到项目所有者头上 - 需另做归属。 + 的配额门照常生效，不要把 followup 也改成准入校验。 +- 费用归属与豁免（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 模型有价，默认 模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL，USD 配额对它们永远不生效（按 0 计）。因此有配额行的用户另受一 @@ -1017,15 +1042,39 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳 时该条回 pending、任务 interrupted，不再排空剩余；建消息时 不查，排队本身不花钱）；Local popen 路径在轮询里每 30 秒查 一次，超限时杀进程组（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`，已花成本照记 - actual，不发终态推送（与取消同款）。诚实边界：Docker/同步 - 路径没有轮询钩子，只有边界检查，单轮内的超支要到轮次结束 - 才停；单轮本身花的钱无法中途按 token 流式截断（中转用量只 - 在回合结束才回传），所以“掐断”粒度是轮次边界 + 长轮询中的 - 30 秒检查，不是 token 级；并发下两个任务同时查到未超限再同 - 时开跑会有少量超支（与限流器同款 best-effort，非精确配额）； - PR 任务（_run_pr_jobs）暂不查配额。 + actual，不发终态推送（与取消同款）。诚实边界：轮询查的是 + 已记账用量，单轮本身花的钱无法中途按 token 流式截断（中转 + 用量只在回合结束才回传），所以“掐断”粒度是轮次边界 + 长轮 + 询中的 30 秒检查，不是 token 级；真 Docker 掐断未实弹（只 + 由假执行器单测覆盖），部署后需实弹；并发下两个任务同时查 + 到未超限再同时开跑会有少量超支（与限流器同款 best-effort， + 非精确配额）——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、 used、remaining，以及当月无价实际 token 数 `unpriced_actual_tokens`——价目表里没有该模型价格的实际 diff --git a/cloud/pr_exec.py b/cloud/pr_exec.py index 9d3c22a..b85fb40 100644 --- a/cloud/pr_exec.py +++ b/cloud/pr_exec.py @@ -48,8 +48,21 @@ def ensure_pr_jobs_table(conn): "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)" + "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() @@ -75,15 +88,21 @@ def find_open_pr_job(conn, task_id): "host": row[3], "status": row[4]} -def enqueue_pr_job(conn, task_id, project_id, host): - """排一个 pr 任务（调用方已做去重）。返回 job id。""" +def enqueue_pr_job(conn, task_id, project_id, host, created_by=None): + """排一个 pr 任务（调用方已做去重）。返回 job id。 + + created_by（G13）：请求 /pr 的用户 id，即费用归属人； + 不传时为 None（老调用方 / 老行），worker 跑它时审计回 + 落到 "system"。server 的 _task_pr 队列路必传调用方。 + """ ensure_pr_jobs_table(conn) job_id = uuid.uuid4().hex now = time.time() conn.execute( "INSERT INTO pr_jobs(id, task_id, project_id, host, status," - " created_at, updated_at) VALUES(?, ?, ?, ?, 'queued', ?, ?)", - (job_id, task_id, project_id, host, now, now)) + " created_at, updated_at, created_by)" + " VALUES(?, ?, ?, ?, 'queued', ?, ?, ?)", + (job_id, task_id, project_id, host, now, now, created_by)) conn.commit() 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）。 user_id：审计 pr_created 的操作人；server 端点传调用方， - worker 的 pr 任务循环传 "system"（与 sweep 同款约定）。 + worker 的 pr 任务循环传 job 的 created_by（G13），老行 + 没有 created_by 时回落 "system"（与 sweep 同款约定）。 token 只在本次调用内存里出现，不进库、不进日志。 """ conn = _db_connect(str(db_path)) diff --git a/cloud/sandbox.py b/cloud/sandbox.py index ebd5a58..f5f54f1 100644 --- a/cloud/sandbox.py +++ b/cloud/sandbox.py @@ -189,11 +189,14 @@ class LocalExecutor(Executor): """ 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.config_path = config_path self.runner = runner if runner is not None else _default_runner 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；传了就覆盖配置。 # worker 的 fallback 靠它把换模型后的第二次尝试真正发到新模型上， @@ -205,6 +208,10 @@ class LocalExecutor(Executor): cmd = [sys.executable, self.host_py_path] if self.config_path is not None: cmd += ["--config", self.config_path] + # G9：顶层全局参数，必须放在子命令之前（调用方在 + # _base_command 之后才追加子命令）。 + if getattr(self, "ignore_host_budget", False): + cmd += ["--ignore-host-budget"] effective = model or self.model if effective: cmd += ["--model", effective] @@ -693,7 +700,7 @@ class DockerExecutor(Executor): user="1000:1000", worktrees_root=None, env=None, env_key="XCM_API_KEY", timeout=None, extra_mounts=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: here = os.path.dirname(os.path.abspath(__file__)) host_py_path = os.path.join(os.path.dirname(here), "host.py") @@ -703,6 +710,8 @@ class DockerExecutor(Executor): # 取构造时传的 model，都没有则不传 --model（容器内 host.py # 用配置文件里的 model）。fallback 时 worker 换任务的 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.image = image self.cpus = str(cpus) @@ -916,6 +925,9 @@ class DockerExecutor(Executor): if self.config_path is not None: cmd += ["--config", "/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 if effective: cmd += ["--model", effective] diff --git a/cloud/server.py b/cloud/server.py index e448c9d..29b41f2 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -6368,7 +6368,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): month)，调用方一律回 402——这是钱用完了，不是请求 太快，所以不用 429。 - 准入点：建任务 / 分支 / 重试 / 续跑，均为权限校验 + 准入点：建任务 / 分支 / 重试 / 续跑 / goal-resume，均为权限校验 通过后、写库之前。P0 安全批后，后续消息与任务间 消息创建时同样查双人配额门（见 _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 不同，旧说法已过时）： - Local popen 路径的 worker 轮询每 30 秒查一次配额， - 超限杀进程组；Docker / 同步路径没有轮询钩子，只有 - 认领与轮次边界检查，单轮内的超支要到轮次结束才停； - 单轮本身的钱无法按 token 流式截断（用量只在回合 - 结束才回传），掐断粒度是轮次边界 + 长轮询中的 + Local popen 与 Docker 路径的 worker 轮询每 30 秒查 + 一次配额，超限杀进程组/容器；同步路径没有轮询钩子， + 只有认领与轮次边界检查，单轮内的超支要到轮次结束 + 才停；单轮本身的钱无法按 token 流式截断（用量只在 + 回合结束才回传），掐断粒度是轮次边界 + 长轮询中的 30 秒检查，不是 token 级。""" quota, used, month = self._quota_status(user_id) if quota is not None and used >= quota: @@ -6395,8 +6407,19 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): 方就是创建者，缺省即调用方；重试 / 续跑时费用记在 原任务创建者名下，必须传 task 的 created_by，否则 维护者替他人重试会用错人的额度（旧 bug：重试按 - 操作人查、费用记创建者名下）。""" - blocked = self._check_quota(quota_user_id or user_id) + 操作人查、费用记创建者名下）。 + + 查询异常时 fail-closed：回 503 quota_check_unavailable + 且不改状态。调用方看到 503 应重试请求本身，而不是 + 认为任务已建好——本函数回 True 时调用方一律直接 + return，不写库。""" + try: + 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: quota, used, month = blocked self._send_json(402, {"error": "quota_exceeded", @@ -6690,6 +6713,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): "status": task["status"], }) 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 任务给该主机； # 本机（或 worktree_host 未记录的老任务）走同步路。 worktree_host = task.get("worktree_host") @@ -6701,7 +6732,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): job_id = existing["id"] else: 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: conn.close() self._send_json(200, { diff --git a/cloud/tests/test_pr_jobs.py b/cloud/tests/test_pr_jobs.py index 9c6a672..0d77378 100644 --- a/cloud/tests/test_pr_jobs.py +++ b/cloud/tests/test_pr_jobs.py @@ -26,7 +26,8 @@ try: from cloud.autoscaler import Autoscaler from cloud.db import connect as _db_connect 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.worker import Worker except ImportError: @@ -105,6 +106,46 @@ class PrJobsTableTestCase(unittest.TestCase): finally: 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): def setUp(self): diff --git a/cloud/tests/test_task_options.py b/cloud/tests/test_task_options.py index 06610a9..d8afa54 100644 --- a/cloud/tests/test_task_options.py +++ b/cloud/tests/test_task_options.py @@ -191,6 +191,39 @@ class TaskOptionsTestCase(unittest.TestCase): "untrusted") 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__": unittest.main() diff --git a/cloud/tests/test_task_quotas.py b/cloud/tests/test_task_quotas.py index fa2ef4a..bf1ab0d 100644 --- a/cloud/tests/test_task_quotas.py +++ b/cloud/tests/test_task_quotas.py @@ -17,14 +17,16 @@ import threading import unittest from datetime import date, datetime, timedelta, timezone from pathlib import Path +from types import SimpleNamespace +from unittest import mock try: 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 sys.path.insert(0, str(Path(__file__).resolve().parents[2])) 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): @@ -355,6 +357,96 @@ class ServerQuotasTestCase(unittest.TestCase): self.assertEqual(s, 200, q) 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__": unittest.main() diff --git a/cloud/tests/test_worker_quota.py b/cloud/tests/test_worker_quota.py index 556abe6..aac3159 100644 --- a/cloud/tests/test_worker_quota.py +++ b/cloud/tests/test_worker_quota.py @@ -18,18 +18,22 @@ import sqlite3 import subprocess import sys import tempfile +import threading import time import unittest from pathlib import Path from types import SimpleNamespace +from unittest import mock try: from cloud.cost import CostStore + from cloud.sandbox import DockerExecutor, LocalExecutor 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.sandbox import DockerExecutor, LocalExecutor from cloud.server import init_db from cloud.worker import Worker @@ -330,6 +334,144 @@ class WorkerQuotaTestCase(unittest.TestCase): task, kind="dispatch", quota_check=lambda: True) 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__": unittest.main() diff --git a/cloud/worker.py b/cloud/worker.py index 4ec20a4..8076283 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -372,7 +372,9 @@ class Worker: upload_dir=None, digest_interval_seconds=3600, 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.host_py_path = host_py_path self.config_path = config_path @@ -385,6 +387,17 @@ class Worker: self.digest_interval_seconds = digest_interval_seconds self.outbox_backoff_seconds = outbox_backoff_seconds 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 # env 只用于 key 存在性检查：默认接进程环境，单测可注入 # 一个普通 dict；worker 永远不把 env 里的值写出。 @@ -416,7 +429,15 @@ class Worker: else: self.executor = LocalExecutor( 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._cancel_requested = set() # 真取消用：process_factory 可注入（单测用它返回一个假长跑 @@ -961,7 +982,73 @@ class Worker: return budget_exceeded(self.usage_turns, self.max_turns, self.usage_cost_usd, self.max_cost_usd) - def _quota_exceeded(self, task, usage_total): + 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。 口径与服务端准入（_quota_status）一致：按任务创建者算， @@ -970,29 +1057,27 @@ class Worker: 另有一道无价 token 上限（与服务端 _check_unpriced_token_cap 同口径）：USD 配额不计量无价 模型，默认模型 terra 无价，若无此上限则配额对默认模型 - 永远不生效；有配额行的用户当月无价 token（input + + 永远不生效；有额度约束的用户（无行按默认额度）当月无价 + token（input + output，已记账 + 本轮）达到上限即停。无价上限只计 source='actual' 的行：估算行（estimated）不是真实花费， 计入上限会把“没花钱”算成“花超了”，且与 cost.py _report 的 unpriced_actual_tokens 口径一致。 - P1 执行面三处收口（旧版三处都是 fail-open，已废）： + 属主与额度（P1 口径，与服务端 _quota_status 一致）： - 属主：任务 created_by 为空时先查父任务（协调员子任务 行可能无 created_by，属主跟父任务走）；仍查不到时 拒绝新回合，不再返回 None 放行。created_by 字面为 - "system" 的系统任务按普通属主算（无行时同样走默认 - 额度，不另设无限额的系统预算）。 - - 无配额行：不再是不限额，按默认配额 - _default_quota_usd_micros()（$10/月）算；只有显式 - NULL 行才是不限额（与 P0 的服务端口径一致，管理员 - 显式设 NULL 的行仍被尊重）。 - - 查库失败：返回 "quota_check_failed: ..."（fail-closed）， - 绝不返回 None 放行。调用方分流：新回合边界（认领门、 - 重试/fallback/后续消息边界）见此串即拒绝新回合并记 - worker_quota_stop 审计；执行中轮询（popen 轮询、 - Docker 配额看门狗）见此串只记警告、不中途杀当前轮 - ——查询抖动不应杀掉已在跑的回合，但下一轮边界仍会 - 拒绝（见 run_one 的 _quota_hit 过滤）。查库失败的 - 审计由边界调用方记，这里只记日志，避免轮询高频写审计。 + "system" 的系统任务按普通属主算，不另设无限额的 + 系统预算（G14 的 system 豁免与 P1 冲突，本次未合入）。 + - 无配额行：按默认配额 _default_quota_usd_micros() + （$10/月）算；只有显式 NULL 行才是不限额。 + 查库失败分层（G7 + P1）：认领门传 on_error="closed"， + 返回 "quota_check_unavailable: ..."（调用方标 + interrupted，可重试/续跑）；默认 "open" 返回 + "quota_check_failed: ..."——新回合边界（直接调用） + 见此串即拒绝，执行中轮询经 _quota_hit 闭包过滤此串、 + 只记警告不杀当前轮。查询抖动不应杀掉已在跑的回合， + 但下一轮边界仍会拒绝。 """ owner = (task or {}).get("created_by") if not owner and (task or {}).get("parent_task_id"): @@ -1017,95 +1102,45 @@ class Worker: return ("quota_exceeded: 任务缺少属主，无法查配额，" "拒绝新回合（请归属真实用户后再排队）") 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: - # P1：无行按默认额度，不是不限额（P0 服务端 - # 已在建号/登录时补行，worker 侧仍可能遇到 - # 老库未补行的用户，这里按默认额度兜底， - # 不写行——写行是服务端的职责，worker 不抢）。 - quota = _default_quota_usd_micros() - else: - quota = row[0] - if quota is 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: - 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 micros is None: - # 无价模型：本轮 token 计入无价上限（USD 口 - # 径里它是 0，不会在上面那道里被拦住）。 - unpriced_used += ( - max(0, int(usage_total.get("input_tokens", 0))) - + max(0, int(usage_total.get("output_tokens", 0)))) - if used >= quota: - return ("quota_exceeded: 本月配额已用完（已用 %d / " - "配额 %d 微美元，%s）" % (used, quota, month)) - cap = unpriced_token_cap() - if unpriced_used >= cap: - return ("quota_exceeded: 本月无价模型 token 已用完" - "（已用 %d / 上限 %d，%s）" - % (unpriced_used, cap, month)) - return None + quota, used, month, unpriced_used = \ + self._quota_limit_and_used(owner) except Exception as exc: - # P1 fail-closed：查库失败不再返回 None。调用方按 - # 前缀分流（见本函数 docstring）：quota_check_failed - # 在新回合边界 = 拒绝，在执行中轮询 = 只警告不杀轮。 + # 查库失败分层（见本函数 docstring）：closed（认领门） + # 回 quota_check_unavailable，调用方标 interrupted； + # open 回 quota_check_failed，新回合边界拒绝，执行中 + # 轮询经 _quota_hit 闭包过滤、只警告不杀轮。 try: self.logger.warning("quota check failed: %s", exc) except Exception: pass + if on_error == "closed": + return "quota_check_unavailable: 配额查询失败，任务未开跑" return ("quota_check_failed: 配额查询失败，拒绝新回合" "（查库异常，非配额用完）") + if quota is None: + return None + 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 micros is None: + # 无价模型：本轮 token 计入无价上限（USD 口 + # 径里它是 0，不会在上面那道里被拦住）。 + unpriced_used += ( + max(0, int(usage_total.get("input_tokens", 0))) + + max(0, int(usage_total.get("output_tokens", 0)))) + if used >= quota: + return ("quota_exceeded: 本月配额已用完（已用 %d / " + "配额 %d 微美元，%s）" % (used, quota, month)) + cap = unpriced_token_cap() + if unpriced_used >= cap: + return ("quota_exceeded: 本月无价模型 token 已用完" + "（已用 %d / 上限 %d，%s）" + % (unpriced_used, cap, month)) + return None def _ready_tasks(self): """当前就绪任务列表（pick_next 的批量版）。 @@ -1511,11 +1546,13 @@ class Worker: 的标记，调用方据此保持 interrupted；deadline_seconds 到期时杀进程组并返回 status='timeout' 的标记，调用方 据此置 failed（超时是失败，不是取消）。quota_check 不为 - 空时每 30 秒查一次配额（轮询本身 50 毫秒一圈，逐圈查会 - 把库打爆），超限时同样杀进程组并返回 status= - 'quota_exceeded' 的内部标记——调用方据此走配额中断收尾， - 绝不能把这个标记原样写进 tasks.status（不在合法状态集 - 里）。Docker/同步路径没有这个轮询钩子，只走边界检查。 + 空时每 quota_poll_interval_seconds（缺省 30 秒）查一次 + 配额（轮询本身 50 毫秒一圈，逐圈查会把库打爆），超限时 + 同样杀进程组并返回 status='quota_exceeded' 的内部标记 + ——调用方据此走配额中断收尾，绝不能把这个标记原样写 + 进 tasks.status（不在合法状态集里）。Docker 路径见 + _docker_call_with_quota_poll（同款轮询）；同步路径没有 + 轮询钩子，只走边界检查。 """ task_id = task["id"] gateway_provider, gateway_env = self._ensure_gateway(task) @@ -1539,6 +1576,10 @@ class Worker: started = time.time() last_quota_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: while proc.poll() is None: if task_id in self._cancel_requested or self._loop_stop: @@ -1567,7 +1608,7 @@ class Worker: self._kill_proc_group(proc) break 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() try: over = quota_check() @@ -2205,6 +2246,80 @@ class Worker: detail += " (" + "; ".join(rejected) + ")" 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): """单次执行尝试（不含认领、key 门、重试、成本与最终审计）。 @@ -2337,12 +2452,16 @@ class Worker: if task.get("_goal_resume_breakpoint"): # G10 断点：Docker 路用无 --goal-text 的 # goal-resume-breakpoint 命令。 - result = self.executor.goal_resume_breakpoint_task( + result = self._docker_call_with_quota_poll( task, - timeout=effective_docker_timeout, - extra_mounts=tenant_mounts, - network=tenant_network, - extra_env=proxy_env) + lambda: self.executor.goal_resume_breakpoint_task( + task, + timeout=effective_docker_timeout, + extra_mounts=tenant_mounts, + network=tenant_network, + extra_env=_merge_env(proxy_env, + gateway_env)), + quota_check) elif task.get("_resume"): # P1 断点续跑传 None（不传 --goal-text， # host 沿用原 goal+verify）；开新目标 @@ -2350,31 +2469,41 @@ class Worker: _resume_text = None if task.get( "_resume_kind") == "breakpoint" \ else task["prompt"] - result = self.executor.goal_resume_task( - task, _resume_text, - timeout=effective_docker_timeout, - extra_mounts=tenant_mounts, - network=tenant_network, - extra_env=_merge_env(proxy_env, - gateway_env), - model_provider=gateway_provider) + result = self._docker_call_with_quota_poll( + task, + lambda: self.executor.goal_resume_task( + task, _resume_text, + timeout=effective_docker_timeout, + extra_mounts=tenant_mounts, + network=tenant_network, + extra_env=_merge_env(proxy_env, + gateway_env), + model_provider=gateway_provider), + quota_check) else: - result = self.executor.goal_task( + result = self._docker_call_with_quota_poll( + task, + lambda: self.executor.goal_task( + task, timeout=effective_docker_timeout, + extra_mounts=tenant_mounts, + network=tenant_network, + extra_env=_merge_env(proxy_env, + gateway_env), + seed_host_dir=seed_host_dir, + model_provider=gateway_provider), + quota_check) + else: + result = self._docker_call_with_quota_poll( + task, + lambda: self.executor.dispatch_task( task, timeout=effective_docker_timeout, extra_mounts=tenant_mounts, network=tenant_network, extra_env=_merge_env(proxy_env, gateway_env), seed_host_dir=seed_host_dir, - model_provider=gateway_provider) - else: - result = self.executor.dispatch_task( - task, timeout=effective_docker_timeout, - extra_mounts=tenant_mounts, - network=tenant_network, - extra_env=_merge_env(proxy_env, gateway_env), - seed_host_dir=seed_host_dir, - model_provider=gateway_provider) + model_provider=gateway_provider), + quota_check) else: # seed 参数只传给 LocalExecutor：CloudExecutor # 基类与测试假执行器的 goal_task/dispatch_task @@ -2434,19 +2563,23 @@ class Worker: result = self.executor.dispatch_task(task) combined = (result.get("stdout", "") 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: if is_docker: - result = self.executor.run_task( - task, timeout=effective_docker_timeout, - extra_mounts=tenant_mounts - if self.tenant_isolation else None, - network=tenant_network - if self.tenant_isolation else None, - extra_env=_merge_env( - proxy_env if self.egress_proxy_port - else None, gateway_env), - model_provider=gateway_provider) + result = self._docker_call_with_quota_poll( + task, + lambda: self.executor.run_task( + task, timeout=effective_docker_timeout, + extra_mounts=tenant_mounts + if self.tenant_isolation else None, + network=tenant_network + if self.tenant_isolation else None, + extra_env=_merge_env( + proxy_env if self.egress_proxy_port + else None, gateway_env), + model_provider=gateway_provider), + quota_check) else: if gateway_provider and not isinstance( self.executor, LocalExecutor): @@ -3092,8 +3225,11 @@ class Worker: # 再开跑——直接标 interrupted + 配额原因，不记成本（还没花 # 钱），续跑消息不消费，顶上来后重试/续跑可再跑。与 key # 门、仓库门同款提前返回形状，区别是状态为 interrupted - #（可续跑）而非 failed。 - quota_error = self._quota_exceeded(task, None) + #（可续跑）而非 failed。G7 分层：认领门用 closed—— + # 查询异常时同样标 interrupted（quota_check_unavailable， + # 可重试/续跑），绝不在“不知道超没超”时开跑；轮询与 + # 轮次边界仍用默认 open，查询抖动不杀在跑任务。 + quota_error = self._quota_exceeded(task, None, on_error="closed") if quota_error is not None: self._set_status(task_id, "interrupted") self._write_execution_result( @@ -3696,8 +3832,11 @@ class Worker: job_id = job["id"] task_id = job["task_id"] try: + # G13：审计操作人用 job 的费用归属人（请求 /pr 的 + # 用户），老行没有 created_by 时回落 "system"。 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) try: finish_pr_job(conn, job_id, "completed", @@ -4121,6 +4260,15 @@ def main(argv=None): "列表里的第一个模型再跑一次，401/400" "绝不 fallback，取消绝不 fallback）") 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("--executor", default="local", choices=["local", "docker", "cloud"], @@ -4219,6 +4367,7 @@ def main(argv=None): args.fallback_models.split(",")] if args.fallback_models else []), max_turns=args.max_turns, + ignore_host_budget=args.ignore_host_budget, executor=_build_cli_executor(args, host_py), task_timeout=args.task_timeout, max_retries=args.max_retries, diff --git a/host.py b/host.py index 81d77d3..6a1eed1 100644 --- a/host.py +++ b/host.py @@ -2057,6 +2057,19 @@ def _load_config_file(path): 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): if not p: return p @@ -2099,6 +2112,12 @@ def main(): "打一行 turn progress 并往链路文件记一条 " "kind=progress；0 表示关闭进度输出；不传时用配置 " "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) p = sub.add_parser("init") @@ -2241,6 +2260,10 @@ def main(): print(f"配置文件已写：{config_path}，请按需改 repo/model 后再用") return # 有效配置（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.effort = args.effort or cfg.get("effort") or "low" args.sandbox = args.sandbox or cfg.get("sandbox") diff --git a/tests/test_host.py b/tests/test_host.py index cb88340..688e9e2 100644 --- a/tests/test_host.py +++ b/tests/test_host.py @@ -14,6 +14,23 @@ sys.path.insert(0, str(REPO_ROOT)) 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): def test_save_state_merges_tasks_and_preserves_disk_tasks(self): with tempfile.TemporaryDirectory() as tmp: