多用户生产级缺口批: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:
1 parent
c1f9313c10
commit
7e081211db
7 files changed
+580
-13
No files matched your search
+21
-3
@@ -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
@@ -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>
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user