多用户生产级缺口批:sweep 鉴权、默认配额、默认 executor、 无价 token 上限

- POST /admin/sweep 改走任一项目 maintainer+ 门槛(与备份管理同口径),否则 403,记 admin_sweep 审计;threshold 下限钳 60 秒
- OAuth 建号 created=True 时写默认配额行(默认 $10/月,CLOUD_DEFAULT_QUOTA_USD_MICROS 可调,INSERT OR IGNORE 幂等);老用户无配额行不受影响
- 建任务未指定且无环境继承时默认 executor=docker(分支继承为空时同理);CLOUD_DEFAULT_EXECUTOR 可覆盖,空字符串恢复旧行为;命令台下拉默认改为 docker
- 无价模型 token 上限:价目表外模型(默认 terra)成本 usd 为 NULL,USD 配额永不生效;有配额行的用户当月无价 token 达 CLOUD_QUOTA_UNPRICED_TOKEN_CAP(默认 2000 万/月)即停,服务端 402、worker 认领/重试/后续消息同口径;报表仍记 unpriced,不折假美元
- 测试:新增 test_multiuser_gaps 11 个 + test_worker_quota 3 个;全量 cloud 668、host 64、deploy 19 全过
This commit is contained in:
host committed 2026-10-27 17:15:09 +08:00
1 parent c1f9313c10
commit 7e081211db
7 files changed
+580 -13

No files matched your search

+21 -3
View File
@@ -88,7 +88,10 @@ client_secret(只走环境变量),不是无认证后门。
心跳即注册)、`POST /tasks/{id}/pr`(已 merged 任务的自动 PR,
默认关闭,缺条件时返回 pr_skipped 而不是失败;见下文“云端派发与
自动 PR”一节)。`POST /tasks` 新增可选 `executor`
(local/docker/cloud,未传表示任意 worker 可认领);
(local/docker/cloud;未传且无环境继承时默认 docker——空
executor(任意 worker 可认领)会让任务落到隔离较弱的 local
worker,是多用户下的隔离缺口;`CLOUD_DEFAULT_EXECUTOR` 可
覆盖,空字符串恢复旧行为);
`POST /projects` 新增可选 `auto_pr`(0/1,默认 0)。
### OAuth 真实浏览器登录(已实现,部署前置仍需用户动手)
@@ -373,7 +376,7 @@ psycopg,写明“需安装 psycopg”),绝不静默回落到 SQLite,
| `POST /tasks/{id}/cancel` | 该项目 `create_task` | `queued` / `running` 改为 `interrupted` 并记审计 `cancel_task`;已是终态(completed / failed / interrupted / merged / merge-failed)返回 `400`;不存在 `404` |
| `POST /tasks/{id}/merge` | 该项目 `approve_merge`(maintainer 以上) | 只同步状态,不执行合并:仅 `completed` 可改为 `merged` 并记审计 `merge_task`,其它状态(含已 merged 的重复调用)返回 `409`;不存在 `404`,越权 `403`,无 key `401` |
| `POST /tasks/{id}/heartbeat` | 该项目 `create_task` | 更新 `last_heartbeat = now`;状态为 `queued` 则改为 `running`;返回任务 |
| `POST /admin/sweep` | 认证即可 | 调 `sweep_stale_tasks(db_path, threshold_seconds=300)`(body 可传 `threshold_seconds` 覆盖默认值);把 `running` 且 `last_heartbeat`(无心跳则用 `updated_at`)早于 `now - threshold` 的任务改为 `interrupted` 并记审计 `sweep_interrupted`;返回 `{"swept": [ids]}` |
| `POST /admin/sweep` | 任一项目 maintainer+(与备份管理同口径),否则 `403`;扫完记审计 `admin_sweep` | 调 `sweep_stale_tasks(db_path, threshold_seconds=300)`(body 可传 `threshold_seconds` 覆盖默认值,下限钳到 60 秒);把 `running` 且 `last_heartbeat`(无心跳则用 `updated_at`)早于 `now - threshold` 的任务改为 `interrupted` 并记审计 `sweep_interrupted`;返回 `{"swept": [ids]}` |
| `POST /admin/backup` | maintainer 以上 | Postgres 时调 `backup_postgres`、SQLite 时调 `backup_db`(body 可传 `out_dir`,否则用 `CLOUD_BACKUP_DIR` 或默认 `cloud-backups`);返回 `{backend, path, size_bytes, task_count}`,绝不返回数据库密码或备份文件内容 |
| `GET /tasks/{id}/cost` | 该任务所属项目 `read` | 成本聚合:`{task_id, project_id, total_tokens, usd_micros, turns, fallback_count, actual_tokens, estimated_tokens, cached_input_tokens, unpriced_actual_tokens}`;真数与估算分开,无价 actual 行进 `unpriced_actual_tokens`;不存在 `404`,越权 `403` |
| `POST /projects/{id}/memory` | 该项目 `create_task` | 写线程摘要 body `{thread_id?, goal, conclusion, files_changed?, open_issues?}`,返回 `201 {id}`,并记审计 `create_memory` |
@@ -956,11 +959,26 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳
- 准入:建任务、会话分支、重试在权限校验通过后、写库之前查
配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、
quota_usd_micros、used_usd_micros),不建任务、不改状态。
老用户与未设配额的用户默认不限额,行为与过去一致。重试按
老用户与未设配额的用户默认不限额,行为与过去一致。但
OAuth 新用户不再是“建号即不限额”:建号(created=True)时
写一行默认配额(`CLOUD_DEFAULT_QUOTA_USD_MICROS`,默认
10_000_000,即 $10/月;`INSERT OR IGNORE` 幂等,不覆盖管理
员后来改过的额度),并记 `quota_default` 审计。重试按
操作人查配额、费用仍记到原任务创建者名下(建任务与分支的
操作人就是创建者,没有这个错位)。系统任务(GitHub webhook
的 /fix-review,created_by 是 'system')不记到任何用户名下、
不占配额;要把它算到项目所有者头上需另做归属。
- 无价模型 token 上限:价目表目前只有两个 sol 模型有价,默认
模型 gpt-5.6-terra 等中转模型无价、成本行 usd 为 NULL,USD
配额对它们永远不生效(按 0 计)。因此有配额行的用户另受一
道体积护栏:当月无价模型 token(input + output,已记账 +
本轮,usd_micros IS NULL 的行)达到
`CLOUD_QUOTA_UNPRICED_TOKEN_CAP`(默认 20_000_000/月)即
停——服务端建任务回 402 `quota_exceeded`(附
unpriced_token_cap、unpriced_tokens_used),worker 侧认领/
重试/后续消息轮次同口径拦截。这道是防滥用的体积上限,不
是计价:成本报表里无价 token 仍进 unpriced_actual_tokens,
绝不折算成假美元;无配额行的用户不受限。
- 执行时加固(worker 侧,`Worker._quota_exceeded`,口径与准入
一致:当月已记账 + 本轮已累计用量,按价目表换算):认领后
开跑前再查一次(排队期间配额被花完/调低时不跑,任务标
+1 -2
View File
@@ -151,9 +151,8 @@
<div class="bar">
<select id="envSelect"><option value="">环境:不指定(走项目仓库)</option></select>
<select id="executor">
<option value="">executor:任意 worker</option>
<option value="docker" selected>docker(默认,多用户下隔离最强)</option>
<option value="local">local</option>
<option value="docker">docker</option>
<option value="cloud">cloud</option>
</select>
<label><input id="goalMode" type="checkbox"> goal 模式</label>
+20
View File
@@ -111,6 +111,26 @@ def price_usd_micros(model, input_tokens, cached_input_tokens,
return int(round(usd * 1000000))
def unpriced_token_cap():
"""Monthly unpriced-token ceiling for users who have a quota row.
Models with no price entry (today: the default gpt-5.6-terra and
the other relay models) record usd as unknown, so a USD quota
alone never binds them -- without this cap a quota user could
burn unlimited unpriced tokens. The cap is a volume guardrail,
NOT a price: cost reports still show those tokens as unpriced,
never converted to fake dollars. Only users with a quota row
are capped (no row still means unlimited, as before). Env
``CLOUD_QUOTA_UNPRICED_TOKEN_CAP`` overrides the default
20_000_000 tokens/month; unparsable values fall back to it.
"""
try:
return max(0, int(os.environ.get(
"CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "20000000")))
except (TypeError, ValueError):
return 20000000
_USAGE_SUM_KEYS = ("input_tokens", "cached_input_tokens",
"cache_write_input_tokens", "output_tokens",
"reasoning_output_tokens", "total_tokens", "completions")
+143 -4
View File
@@ -58,7 +58,7 @@ from urllib.parse import parse_qs, quote, urlparse
from cloud.auth import AuthStore
from cloud.coordinator import parse_parts as _coord_parse_parts
from cloud.coordinator import validate_parts as _coord_validate_parts
from cloud.cost import CostStore
from cloud.cost import CostStore, unpriced_token_cap
from cloud.db import connect as _db_connect
from cloud.db import create_database as _create_database
from cloud.db import init_postgres_schema as _init_postgres_schema
@@ -936,6 +936,61 @@ def _ensure_quotas_table(conn):
"updated_at DOUBLE PRECISION)")
def _default_executor():
"""新任务未指定执行器时的默认值:docker。
空 executor(任意 worker 可认领)是多用户下的隔离缺口:
任务可能落到隔离较弱的 local worker 上。默认 docker 后,
只有显式 local/cloud 或系统创建(webhook/修复任务,executor
为 NULL)才走别的路。环境变量 CLOUD_DEFAULT_EXECUTOR 可
覆盖:设为空字符串即恢复旧行为(None,任意 worker),
供单机运维在没有 docker worker 时兜底;非法值按 docker 计。
"""
raw = os.environ.get("CLOUD_DEFAULT_EXECUTOR", "docker")
if raw is None:
return "docker"
raw = raw.strip()
if raw == "":
return None
if raw in ("local", "docker", "cloud"):
return raw
return "docker"
def _default_quota_usd_micros():
"""OAuth 新用户建号时写入的默认配额(微美元)。
无配额行 = 不限额,所以 OAuth 自动建号过去等于无限额入驻。
新号现在建号即写一行默认配额(老用户无配额行不受影响,
行为不变)。环境变量 CLOUD_DEFAULT_QUOTA_USD_MICROS 可
覆盖,默认 10_000_000($10/月);非法值按默认计。
"""
try:
return max(0, int(os.environ.get(
"CLOUD_DEFAULT_QUOTA_USD_MICROS", "10000000")))
except (TypeError, ValueError):
return 10000000
def _apply_default_quota(conn, user_id):
"""给新用户写默认配额行(已存在则不动,幂等)。
只在 OAuth 建号(created=True)时调:INSERT OR IGNORE
保证重复回调不覆盖管理员后来改过的额度。返回 True 表示
本次写了新行(调用方据此记 quota_default 审计)。
"""
_ensure_quotas_table(conn)
cur = conn.execute(
"INSERT OR IGNORE INTO user_quotas(user_id, quota_usd_micros,"
" updated_at) VALUES(?, ?, ?)",
(user_id, _default_quota_usd_micros(), time.time()))
try:
conn.commit()
except Exception:
pass
return cur.rowcount == 1
_ENV_BRANCH_RE = re.compile(r"^[A-Za-z0-9_.\-/]+$")
@@ -1465,7 +1520,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
elif path == "/tasks" and method == "GET":
self._list_tasks(user_id, query)
elif path == "/admin/sweep" and method == "POST":
self._admin_sweep()
self._admin_sweep(user_id)
elif path == "/admin/backup" and method == "POST":
self._admin_backup(user_id)
elif path == "/admin/quotas" and method == "GET":
@@ -2582,6 +2637,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
return
if executor is None and env["executor"]:
executor = env["executor"]
if executor is None:
# 未指定且无环境继承时默认 docker:空 executor(任意
# worker 可认领)会让任务落到隔离较弱的 local worker,
# 是多用户下的隔离缺口;显式 local/cloud 不受影响。
executor = _default_executor()
# 附件绑定(upload_id 由 POST /uploads 发出):包必须存在、
# 归属当前用户、未使用、24 小时内、暂存目录还在;cloud
# 执行器没有 worktree 可 seed,直接拒绝。校验在 INSERT
@@ -3288,7 +3348,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
" VALUES(?, ?, ?, 'queued', NULL, '[]', ?, ?, ?, ?,"
" NULL, ?, NULL, NULL, ?, ?, ?, ?, ?, ?)",
(new_id, task["project_id"], message, task["priority"],
user_id, now, now, task["executor"],
user_id, now, now,
task["executor"] or _default_executor(),
task["goal_max_rounds"] or 3, task["model"],
task["effort"], task["approval"], task["plan"] or 0,
task_id),
@@ -3809,6 +3870,56 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
"remaining_usd_micros": (quota - used)
if quota is not None else None}
def _quota_unpriced_tokens(self, user_id, month):
"""当月无价模型 token 数(usd_micros IS NULL 的成本行,
input + output 之和;cached 已含在 input 里,不重复计)。
与 _quota_status 同款 CostStore 懒建表做法,全新库上
cost_records 不存在时先建表再查。无配额行时调用方不
应以此挡人(见 _check_unpriced_token_cap)。
"""
conn = self._db()
try:
_cost_store = CostStore(self.server.db_path)
try:
_cost_store._conn.close()
except Exception:
pass
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",
(user_id, month + "%")).fetchone()
return int(row[0] if row is not None else 0)
finally:
conn.close()
def _check_unpriced_token_cap(self, user_id):
"""无价 token 上限准入:超限返回 (cap, used, month),否则 None。
只有有配额行的用户才受限(无配额行 = 不限额,与 USD
口径一致);USD 配额本身不计量无价模型,所以这道是
补 USD 口径的洞,不是第二套计费。"""
month = self._quota_month()
conn = self._db()
try:
_ensure_quotas_table(conn)
conn.commit()
row = conn.execute(
"SELECT quota_usd_micros FROM user_quotas"
" WHERE user_id = ?",
(user_id,)).fetchone()
if row is None:
return None
finally:
conn.close()
used = self._quota_unpriced_tokens(user_id, month)
cap = unpriced_token_cap()
if used >= cap:
return cap, used, month
return None
def _check_quota(self, user_id):
"""配额准入:放行返回 None;超限返回 (quota, used,
month),调用方一律回 402——这是钱用完了,不是请求
@@ -3829,6 +3940,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
"quota_usd_micros": quota,
"used_usd_micros": used})
return True
capped = self._check_unpriced_token_cap(user_id)
if capped is not None:
cap, used, month = capped
self._send_json(402, {"error": "quota_exceeded",
"month": month,
"unpriced_token_cap": cap,
"unpriced_tokens_used": used})
return True
return False
def _is_global_maintainer(self, user_id):
@@ -4115,7 +4234,19 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
self._send_json(200, task)
# -- 管理 ----------------------------------------------------------
def _admin_sweep(self):
def _admin_sweep(self, user_id):
"""POST /admin/sweep:把超阈值的陈旧 running 任务标 interrupted。
权限:与 _admin_backup 同口径——调用方须在任一项目里
是 maintainer 以上,否则 403。过去这里没有任何权限检查,
任意持 key 用户都能跨用户改任务状态。threshold_seconds
下限钳到 60 秒:0 会把所有 running 任务一口气扫掉,
只能是管理员的显式紧急操作才需要那么小的值,而 60 秒
以下的“陈旧”本就没有意义。扫完记 admin_sweep 审计。
"""
if not self._is_global_maintainer(user_id):
self._send_json(403, {"error": "forbidden"})
return
body, err = self._read_json()
if err:
self._send_json(400, {"error": err})
@@ -4127,7 +4258,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
except (TypeError, ValueError):
self._send_json(400, {"error": "threshold_seconds must be an integer"})
return
threshold = max(60, threshold)
swept = sweep_stale_tasks(self.server.db_path, threshold)
with _auth_store(self.server.db_path) as store:
store.audit(user_id, "", "admin_sweep",
"threshold=%d swept=%d" % (threshold, len(swept)))
self._send_json(200, {"swept": swept})
def _admin_backup(self, user_id):
@@ -4366,6 +4501,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
auth, provider, identity)
auth.audit(local_id, "", "oauth_login",
"%s:%s" % (provider, identity.get("external_id")))
if created and _apply_default_quota(auth.conn, local_id):
auth.audit(local_id, "", "quota_default",
"quota_usd_micros=%d"
% _default_quota_usd_micros())
if not created:
auth.audit(local_id, "", "api_key_issued", provider)
return local_id, created, api_key_plain, rotated, identity, False
+279
View File
@@ -0,0 +1,279 @@
"""多用户生产级缺口批单测:sweep 鉴权、默认执行器、新用户默认配额、
无价模型 token 上限(服务端准入)。
只用本地临时端口 + 临时 SQLite,不跑真实模型,不联网。
环境变量(默认执行器、默认配额额、token 上限)是服务端每次
请求现读的,测试里用 try/finally 改回,不影响其他测试。
"""
import http.client
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from datetime import datetime, timezone
from pathlib import Path
try:
from cloud.auth import AuthStore
from cloud.cost import CostStore
from cloud.server import _apply_default_quota, _default_quota_usd_micros
from cloud.server import 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.cost import CostStore
from cloud.server import _apply_default_quota, _default_quota_usd_micros
from cloud.server import create_server
def _req(port, method, path, body=None, key=None):
conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10)
hdrs = {}
payload = None
if body is not None:
payload = json.dumps(body).encode("utf-8")
hdrs["Content-Type"] = "application/json"
if key is not None:
hdrs["X-API-Key"] = key
try:
conn.request(method, path, body=payload, headers=hdrs)
resp = conn.getresponse()
data = resp.read()
finally:
conn.close()
try:
parsed = json.loads(data.decode("utf-8")) if data else None
except ValueError:
parsed = None
return resp.status, parsed
def _utc_today():
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
def _setenv(name, value):
old = os.environ.get(name)
if value is None:
os.environ.pop(name, None)
else:
os.environ[name] = value
return old
class MultiuserGapsTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.db_path = os.path.join(self.tmp.name, "control.db")
self.auth = AuthStore(self.db_path)
self.owner_id, self.owner_key = self.auth.create_user("owner")
self.viewer_id, self.viewer_key = self.auth.create_user("viewer")
self.outsider_id, self.outsider_key = self.auth.create_user("outsider")
self.server = create_server(self.db_path, "127.0.0.1", 0)
self.port = self.server.server_address[1]
self.thread = threading.Thread(
target=self.server.serve_forever, daemon=True)
self.thread.start()
s, proj = _req(self.port, "POST", "/projects",
{"name": "gproj"}, self.owner_key)
self.assertEqual(s, 201, proj)
self.project_id = proj["id"]
self.auth.set_member(
self.project_id, self.owner_id, self.viewer_id, "viewer")
def tearDown(self):
self.server.shutdown()
self.server.server_close()
self.thread.join(timeout=5)
try:
self.auth.conn.close()
except Exception:
pass
self.tmp.cleanup()
# -- helpers -------------------------------------------------------
def _create_task(self, key=None, body_extra=None):
body = {"project_id": self.project_id, "prompt": "gap probe"}
if body_extra:
body.update(body_extra)
return _req(self.port, "POST", "/tasks", body,
key if key is not None else self.owner_key)
def _insert_quota(self, user_id, micros):
conn = sqlite3.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.execute(
"INSERT OR REPLACE INTO user_quotas"
"(user_id, quota_usd_micros, updated_at) VALUES(?, ?, ?)",
(user_id, micros, time.time()))
conn.commit()
finally:
conn.close()
def _insert_unpriced_cost(self, task_id, input_tokens, output_tokens):
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"INSERT INTO tasks(id, project_id, prompt, status,"
" idempotency_key, depends_on, priority, created_by,"
" created_at, updated_at, last_heartbeat)"
" VALUES(?, ?, ?, 'completed', NULL, '[]', 0, ?, ?, ?, NULL)",
(task_id, self.project_id, "donor", self.owner_id,
time.time(), time.time()))
conn.commit()
finally:
conn.close()
# CostStore.record 建表并按 UTC 今天记 day,与服务端查询口径一致。
store = CostStore(self.db_path)
try:
store.record(task_id, self.project_id, "gpt-5.6-terra",
input_tokens=input_tokens,
output_tokens=output_tokens,
usd_micros=None, source="actual")
finally:
try:
store._conn.close()
except Exception:
pass
def _make_running_task(self):
s, task = self._create_task()
self.assertEqual(s, 201, task)
s, _ = _req(self.port, "POST",
"/tasks/%s/heartbeat" % task["id"], {},
self.owner_key)
self.assertEqual(s, 200)
return task["id"]
# -- sweep 鉴权 -----------------------------------------------------
def test_sweep_viewer_and_outsider_forbidden(self):
for key in (self.viewer_key, self.outsider_key):
s, body = _req(self.port, "POST", "/admin/sweep", {},
key)
self.assertEqual(s, 403, body)
def test_sweep_owner_maintainer_ok_and_sweeps_stale(self):
tid = self._make_running_task()
conn = sqlite3.connect(self.db_path)
try:
conn.execute(
"UPDATE tasks SET last_heartbeat = ? WHERE id = ?",
(time.time() - 3600, tid))
conn.commit()
finally:
conn.close()
s, body = _req(self.port, "POST", "/admin/sweep",
{"threshold_seconds": 300}, self.owner_key)
self.assertEqual(s, 200, body)
self.assertIn(tid, body["swept"])
def test_sweep_threshold_clamped_to_60(self):
tid = self._make_running_task()
s, body = _req(self.port, "POST", "/admin/sweep",
{"threshold_seconds": 0}, self.owner_key)
self.assertEqual(s, 200, body)
self.assertNotIn(tid, body["swept"])
# -- 默认执行器 ------------------------------------------------------
def test_default_executor_is_docker(self):
s, task = self._create_task()
self.assertEqual(s, 201, task)
self.assertEqual(task["executor"], "docker")
def test_explicit_executor_kept(self):
s, task = self._create_task(body_extra={"executor": "local"})
self.assertEqual(s, 201, task)
self.assertEqual(task["executor"], "local")
def test_default_executor_env_override(self):
old = _setenv("CLOUD_DEFAULT_EXECUTOR", "local")
try:
s, task = self._create_task()
self.assertEqual(s, 201, task)
self.assertEqual(task["executor"], "local")
finally:
_setenv("CLOUD_DEFAULT_EXECUTOR", old)
old = _setenv("CLOUD_DEFAULT_EXECUTOR", "")
try:
s, task = self._create_task()
self.assertEqual(s, 201, task)
self.assertIsNone(task["executor"])
finally:
_setenv("CLOUD_DEFAULT_EXECUTOR", old)
# -- 新用户默认配额 ---------------------------------------------------
def test_apply_default_quota_idempotent(self):
conn = sqlite3.connect(self.db_path)
try:
self.assertTrue(_apply_default_quota(conn, "newbie"))
row = conn.execute(
"SELECT quota_usd_micros FROM user_quotas"
" WHERE user_id = 'newbie'").fetchone()
self.assertEqual(row[0], _default_quota_usd_micros())
# 第二次不再写,不覆盖管理员后来改过的额度。
conn.execute(
"UPDATE user_quotas SET quota_usd_micros = 12345"
" WHERE user_id = 'newbie'")
conn.commit()
self.assertFalse(_apply_default_quota(conn, "newbie"))
row = conn.execute(
"SELECT quota_usd_micros FROM user_quotas"
" WHERE user_id = 'newbie'").fetchone()
self.assertEqual(row[0], 12345)
finally:
conn.close()
def test_default_quota_env_override(self):
old = _setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", "5000")
try:
self.assertEqual(_default_quota_usd_micros(), 5000)
finally:
_setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", old)
# -- 无价 token 上限(服务端准入)--------------------------------------
def test_unpriced_token_cap_blocks_create(self):
self._insert_quota(self.owner_id, 10 ** 12) # USD 很大,不挡
for i in range(7): # 7 * 150 = 1050 token
self._insert_unpriced_cost("donor-%d" % i, 100, 50)
old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000")
try:
s, body = self._create_task()
self.assertEqual(s, 402, body)
self.assertEqual(body["error"], "quota_exceeded")
self.assertEqual(body["unpriced_tokens_used"], 1050)
finally:
_setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old)
def test_unpriced_token_cap_under_cap_passes(self):
self._insert_quota(self.owner_id, 10 ** 12)
for i in range(6): # 6 * 150 = 900 < 1000
self._insert_unpriced_cost("donor-%d" % i, 100, 50)
old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000")
try:
s, task = self._create_task()
self.assertEqual(s, 201, task)
finally:
_setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old)
def test_unpriced_tokens_unlimited_without_quota_row(self):
for i in range(7):
self._insert_unpriced_cost("donor-%d" % i, 100, 50)
old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000")
try:
s, task = self._create_task()
self.assertEqual(s, 201, task)
finally:
_setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old)
if __name__ == "__main__":
unittest.main()
+87
View File
@@ -80,6 +80,26 @@ def _insert_donor_cost(db_path, micros):
pass
def _insert_unpriced_cost(db_path, tokens):
"""给 tester 记一笔无价成本:已完成任务 + usd_micros NULL 的行。
tokens 拆成 input/output 两半记(input + output = tokens,
与服务端/ worker 的无价 token 口径一致)。
"""
_insert_task(db_path, "donor-unpriced", status="completed")
store = CostStore(db_path)
try:
store.record("donor-unpriced", "p1", "gpt-5.6-terra",
input_tokens=tokens // 2,
output_tokens=tokens - tokens // 2,
usd_micros=None, source="actual")
finally:
try:
store._conn.close()
except Exception:
pass
def _insert_followup(db_path, fu_id, task_id, body="follow up"):
now = time.time()
conn = sqlite3.connect(db_path)
@@ -229,6 +249,73 @@ class WorkerQuotaTestCase(unittest.TestCase):
result = worker.run_one()
self.assertEqual(result["status"], "completed")
# -- 无价 token 上限 -------------------------------------------------
def test_unpriced_token_cap_blocks_claim(self):
_insert_quota(self.db_path, "tester", 10 ** 12) # USD 很大,不挡
_insert_unpriced_cost(self.db_path, 150)
_insert_task(self.db_path, "T1")
worker = self._worker(_usage_runner(
{"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
self.calls))
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
try:
result = worker.run_one()
finally:
if old is None:
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
else:
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
self.assertEqual(result["status"], "interrupted")
self.assertIn("无价模型 token", result["error"])
self.assertEqual(self.calls, [])
def test_unpriced_round_usage_counts_toward_cap(self):
_insert_quota(self.db_path, "tester", 10 ** 12)
_insert_unpriced_cost(self.db_path, 90)
_insert_task(self.db_path, "T1")
_insert_followup(self.db_path, "fu1", "T1")
worker = Worker(self.db_path, "/fake/host.py", runner=_usage_runner(
{"input_tokens": 50, "cached_input_tokens": 0,
"output_tokens": 0, "total_tokens": 50},
self.calls), env=self.env, model="gpt-5.6-terra",
retry_backoff_seconds=0)
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
try:
result = worker.run_one()
finally:
if old is None:
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
else:
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
# 首轮无价 50 token:90 + 50 = 140 >= 100,排空时停。
self.assertEqual(result["status"], "interrupted")
self.assertIn("quota_exceeded", result["error"])
self.assertEqual(_followup_status(self.db_path, "fu1"), "pending")
def test_priced_usage_ignores_token_cap(self):
_insert_quota(self.db_path, "tester", 10 ** 12)
_insert_unpriced_cost(self.db_path, 90)
_insert_task(self.db_path, "T1")
_insert_followup(self.db_path, "fu1", "T1")
worker = self._worker(_usage_runner(
{"input_tokens": 100, "cached_input_tokens": 0,
"output_tokens": 10, "total_tokens": 110},
self.calls))
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
try:
result = worker.run_one()
finally:
if old is None:
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
else:
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
# 有价模型用量不计入无价上限:90 < 100,正常跑完。
self.assertEqual(result["status"], "completed")
self.assertEqual(_followup_status(self.db_path, "fu1"), "done")
# -- 轮询掐断 ------------------------------------------------------
def test_poll_loop_quota_kill(self):
_insert_task(self.db_path, "T1")
+29 -4
View File
@@ -139,7 +139,8 @@ from cloud.coordinator import DEFAULT_POLL_INTERVAL as _COORD_POLL_INTERVAL
from cloud.coordinator import DEFAULT_TIMEOUT_SECONDS as _COORD_TIMEOUT
from cloud.coordinator import ensure_columns as _coord_ensure_columns
from cloud.coordinator import parse_parts as _coord_parse_parts
from cloud.cost import CostStore, estimate_tokens, price_usd_micros, sum_usage
from cloud.cost import CostStore, estimate_tokens, price_usd_micros, \
sum_usage, unpriced_token_cap
from cloud.db import connect as _db_connect
from cloud.logging_config import get_logger, log_event, setup_logging
from cloud.models import ModelRouter, is_fallbackable, record_fallback
@@ -832,9 +833,14 @@ class Worker:
口径与服务端准入(_quota_status)一致:按任务创建者算,
当月(UTC)已记账用量 + 本轮已累计用量(按价目表换算,
无价模型按 0 计,与服务端 SUM(usd NULL 不计)一致)。
无配额行 = 不限额。查询本身异常时按不限额处理(fail-open
并记一条警告):配额是第二道网,绝不能因查询抖动把全部
任务挡在门外;服务端建任务/分支/重试时的准入仍是第一道。
另有一道无价 token 上限(与服务端
_check_unpriced_token_cap 同口径):USD 配额不计量无价
模型,默认模型 terra 无价,若无此上限则配额对默认模型
永远不生效;有配额行的用户当月无价 token(input +
output,已记账 + 本轮)达到上限即停。无配额行 = 不限额。
查询本身异常时按不限额处理(fail-open 并记一条警告):
配额是第二道网,绝不能因查询抖动把全部任务挡在门外;
服务端建任务/分支/重试时的准入仍是第一道。
"""
owner = (task or {}).get("created_by")
if not owner:
@@ -873,6 +879,14 @@ class Worker:
" 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",
(owner, month + "%")).fetchone()
unpriced_used = unpriced_row[0] if unpriced_row else 0
finally:
conn.close()
if usage_total:
@@ -882,9 +896,20 @@ class Worker:
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
except Exception as exc:
try: