Merge pull request #86 from openqj/feat-multiuser-hardening
多用户生产级缺口批:sweep 鉴权、默认配额、默认 executor、 无价 token 上限
This commit is contained in:
9 files changed
+627
-14
No files matched your search
+21
-3
@@ -88,7 +88,10 @@ client_secret(只走环境变量),不是无认证后门。
|
|||||||
心跳即注册)、`POST /tasks/{id}/pr`(已 merged 任务的自动 PR,
|
心跳即注册)、`POST /tasks/{id}/pr`(已 merged 任务的自动 PR,
|
||||||
默认关闭,缺条件时返回 pr_skipped 而不是失败;见下文“云端派发与
|
默认关闭,缺条件时返回 pr_skipped 而不是失败;见下文“云端派发与
|
||||||
自动 PR”一节)。`POST /tasks` 新增可选 `executor`
|
自动 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)。
|
`POST /projects` 新增可选 `auto_pr`(0/1,默认 0)。
|
||||||
|
|
||||||
### OAuth 真实浏览器登录(已实现,部署前置仍需用户动手)
|
### 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}/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}/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 /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}`,绝不返回数据库密码或备份文件内容 |
|
| `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` |
|
| `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` |
|
| `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、
|
配额,本月用量 >= 配额时回 402 `quota_exceeded`(附 month、
|
||||||
quota_usd_micros、used_usd_micros),不建任务、不改状态。
|
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
|
操作人就是创建者,没有这个错位)。系统任务(GitHub webhook
|
||||||
的 /fix-review,created_by 是 'system')不记到任何用户名下、
|
的 /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`,口径与准入
|
- 执行时加固(worker 侧,`Worker._quota_exceeded`,口径与准入
|
||||||
一致:当月已记账 + 本轮已累计用量,按价目表换算):认领后
|
一致:当月已记账 + 本轮已累计用量,按价目表换算):认领后
|
||||||
开跑前再查一次(排队期间配额被花完/调低时不跑,任务标
|
开跑前再查一次(排队期间配额被花完/调低时不跑,任务标
|
||||||
|
|||||||
+1
-2
@@ -151,9 +151,8 @@
|
|||||||
<div class="bar">
|
<div class="bar">
|
||||||
<select id="envSelect"><option value="">环境:不指定(走项目仓库)</option></select>
|
<select id="envSelect"><option value="">环境:不指定(走项目仓库)</option></select>
|
||||||
<select id="executor">
|
<select id="executor">
|
||||||
<option value="">executor:任意 worker</option>
|
<option value="docker" selected>docker(默认,多用户下隔离最强)</option>
|
||||||
<option value="local">local</option>
|
<option value="local">local</option>
|
||||||
<option value="docker">docker</option>
|
|
||||||
<option value="cloud">cloud</option>
|
<option value="cloud">cloud</option>
|
||||||
</select>
|
</select>
|
||||||
<label><input id="goalMode" type="checkbox"> goal 模式</label>
|
<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))
|
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",
|
_USAGE_SUM_KEYS = ("input_tokens", "cached_input_tokens",
|
||||||
"cache_write_input_tokens", "output_tokens",
|
"cache_write_input_tokens", "output_tokens",
|
||||||
"reasoning_output_tokens", "total_tokens", "completions")
|
"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.auth import AuthStore
|
||||||
from cloud.coordinator import parse_parts as _coord_parse_parts
|
from cloud.coordinator import parse_parts as _coord_parse_parts
|
||||||
from cloud.coordinator import validate_parts as _coord_validate_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 connect as _db_connect
|
||||||
from cloud.db import create_database as _create_database
|
from cloud.db import create_database as _create_database
|
||||||
from cloud.db import init_postgres_schema as _init_postgres_schema
|
from cloud.db import init_postgres_schema as _init_postgres_schema
|
||||||
@@ -936,6 +936,61 @@ def _ensure_quotas_table(conn):
|
|||||||
"updated_at DOUBLE PRECISION)")
|
"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_.\-/]+$")
|
_ENV_BRANCH_RE = re.compile(r"^[A-Za-z0-9_.\-/]+$")
|
||||||
|
|
||||||
|
|
||||||
@@ -1465,7 +1520,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
elif path == "/tasks" and method == "GET":
|
elif path == "/tasks" and method == "GET":
|
||||||
self._list_tasks(user_id, query)
|
self._list_tasks(user_id, query)
|
||||||
elif path == "/admin/sweep" and method == "POST":
|
elif path == "/admin/sweep" and method == "POST":
|
||||||
self._admin_sweep()
|
self._admin_sweep(user_id)
|
||||||
elif path == "/admin/backup" and method == "POST":
|
elif path == "/admin/backup" and method == "POST":
|
||||||
self._admin_backup(user_id)
|
self._admin_backup(user_id)
|
||||||
elif path == "/admin/quotas" and method == "GET":
|
elif path == "/admin/quotas" and method == "GET":
|
||||||
@@ -2582,6 +2637,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
return
|
return
|
||||||
if executor is None and env["executor"]:
|
if executor is None and env["executor"]:
|
||||||
executor = env["executor"]
|
executor = env["executor"]
|
||||||
|
if executor is None:
|
||||||
|
# 未指定且无环境继承时默认 docker:空 executor(任意
|
||||||
|
# worker 可认领)会让任务落到隔离较弱的 local worker,
|
||||||
|
# 是多用户下的隔离缺口;显式 local/cloud 不受影响。
|
||||||
|
executor = _default_executor()
|
||||||
# 附件绑定(upload_id 由 POST /uploads 发出):包必须存在、
|
# 附件绑定(upload_id 由 POST /uploads 发出):包必须存在、
|
||||||
# 归属当前用户、未使用、24 小时内、暂存目录还在;cloud
|
# 归属当前用户、未使用、24 小时内、暂存目录还在;cloud
|
||||||
# 执行器没有 worktree 可 seed,直接拒绝。校验在 INSERT
|
# 执行器没有 worktree 可 seed,直接拒绝。校验在 INSERT
|
||||||
@@ -3288,7 +3348,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
" VALUES(?, ?, ?, 'queued', NULL, '[]', ?, ?, ?, ?,"
|
" VALUES(?, ?, ?, 'queued', NULL, '[]', ?, ?, ?, ?,"
|
||||||
" NULL, ?, NULL, NULL, ?, ?, ?, ?, ?, ?)",
|
" NULL, ?, NULL, NULL, ?, ?, ?, ?, ?, ?)",
|
||||||
(new_id, task["project_id"], message, task["priority"],
|
(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["goal_max_rounds"] or 3, task["model"],
|
||||||
task["effort"], task["approval"], task["plan"] or 0,
|
task["effort"], task["approval"], task["plan"] or 0,
|
||||||
task_id),
|
task_id),
|
||||||
@@ -3809,6 +3870,56 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
"remaining_usd_micros": (quota - used)
|
"remaining_usd_micros": (quota - used)
|
||||||
if quota is not None else None}
|
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):
|
def _check_quota(self, user_id):
|
||||||
"""配额准入:放行返回 None;超限返回 (quota, used,
|
"""配额准入:放行返回 None;超限返回 (quota, used,
|
||||||
month),调用方一律回 402——这是钱用完了,不是请求
|
month),调用方一律回 402——这是钱用完了,不是请求
|
||||||
@@ -3829,6 +3940,14 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
"quota_usd_micros": quota,
|
"quota_usd_micros": quota,
|
||||||
"used_usd_micros": used})
|
"used_usd_micros": used})
|
||||||
return True
|
return True
|
||||||
|
capped = self._check_unpriced_token_cap(user_id)
|
||||||
|
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
|
return False
|
||||||
|
|
||||||
def _is_global_maintainer(self, user_id):
|
def _is_global_maintainer(self, user_id):
|
||||||
@@ -4115,7 +4234,19 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
self._send_json(200, task)
|
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()
|
body, err = self._read_json()
|
||||||
if err:
|
if err:
|
||||||
self._send_json(400, {"error": err})
|
self._send_json(400, {"error": err})
|
||||||
@@ -4127,7 +4258,11 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
except (TypeError, ValueError):
|
except (TypeError, ValueError):
|
||||||
self._send_json(400, {"error": "threshold_seconds must be an integer"})
|
self._send_json(400, {"error": "threshold_seconds must be an integer"})
|
||||||
return
|
return
|
||||||
|
threshold = max(60, threshold)
|
||||||
swept = sweep_stale_tasks(self.server.db_path, 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})
|
self._send_json(200, {"swept": swept})
|
||||||
|
|
||||||
def _admin_backup(self, user_id):
|
def _admin_backup(self, user_id):
|
||||||
@@ -4366,6 +4501,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
|||||||
auth, provider, identity)
|
auth, provider, identity)
|
||||||
auth.audit(local_id, "", "oauth_login",
|
auth.audit(local_id, "", "oauth_login",
|
||||||
"%s:%s" % (provider, identity.get("external_id")))
|
"%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:
|
if not created:
|
||||||
auth.audit(local_id, "", "api_key_issued", provider)
|
auth.audit(local_id, "", "api_key_issued", provider)
|
||||||
return local_id, created, api_key_plain, rotated, identity, False
|
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
|
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"):
|
def _insert_followup(db_path, fu_id, task_id, body="follow up"):
|
||||||
now = time.time()
|
now = time.time()
|
||||||
conn = sqlite3.connect(db_path)
|
conn = sqlite3.connect(db_path)
|
||||||
@@ -229,6 +249,73 @@ class WorkerQuotaTestCase(unittest.TestCase):
|
|||||||
result = worker.run_one()
|
result = worker.run_one()
|
||||||
self.assertEqual(result["status"], "completed")
|
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):
|
def test_poll_loop_quota_kill(self):
|
||||||
_insert_task(self.db_path, "T1")
|
_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 DEFAULT_TIMEOUT_SECONDS as _COORD_TIMEOUT
|
||||||
from cloud.coordinator import ensure_columns as _coord_ensure_columns
|
from cloud.coordinator import ensure_columns as _coord_ensure_columns
|
||||||
from cloud.coordinator import parse_parts as _coord_parse_parts
|
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.db import connect as _db_connect
|
||||||
from cloud.logging_config import get_logger, log_event, setup_logging
|
from cloud.logging_config import get_logger, log_event, setup_logging
|
||||||
from cloud.models import ModelRouter, is_fallbackable, record_fallback
|
from cloud.models import ModelRouter, is_fallbackable, record_fallback
|
||||||
@@ -832,9 +833,14 @@ class Worker:
|
|||||||
口径与服务端准入(_quota_status)一致:按任务创建者算,
|
口径与服务端准入(_quota_status)一致:按任务创建者算,
|
||||||
当月(UTC)已记账用量 + 本轮已累计用量(按价目表换算,
|
当月(UTC)已记账用量 + 本轮已累计用量(按价目表换算,
|
||||||
无价模型按 0 计,与服务端 SUM(usd NULL 不计)一致)。
|
无价模型按 0 计,与服务端 SUM(usd NULL 不计)一致)。
|
||||||
无配额行 = 不限额。查询本身异常时按不限额处理(fail-open
|
另有一道无价 token 上限(与服务端
|
||||||
并记一条警告):配额是第二道网,绝不能因查询抖动把全部
|
_check_unpriced_token_cap 同口径):USD 配额不计量无价
|
||||||
任务挡在门外;服务端建任务/分支/重试时的准入仍是第一道。
|
模型,默认模型 terra 无价,若无此上限则配额对默认模型
|
||||||
|
永远不生效;有配额行的用户当月无价 token(input +
|
||||||
|
output,已记账 + 本轮)达到上限即停。无配额行 = 不限额。
|
||||||
|
查询本身异常时按不限额处理(fail-open 并记一条警告):
|
||||||
|
配额是第二道网,绝不能因查询抖动把全部任务挡在门外;
|
||||||
|
服务端建任务/分支/重试时的准入仍是第一道。
|
||||||
"""
|
"""
|
||||||
owner = (task or {}).get("created_by")
|
owner = (task or {}).get("created_by")
|
||||||
if not owner:
|
if not owner:
|
||||||
@@ -873,6 +879,14 @@ class Worker:
|
|||||||
" WHERE t.created_by = ? AND c.day LIKE ?",
|
" WHERE t.created_by = ? AND c.day LIKE ?",
|
||||||
(owner, month + "%")).fetchone()
|
(owner, month + "%")).fetchone()
|
||||||
used = used_row[0] if used_row else 0
|
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:
|
finally:
|
||||||
conn.close()
|
conn.close()
|
||||||
if usage_total:
|
if usage_total:
|
||||||
@@ -882,9 +896,20 @@ class Worker:
|
|||||||
usage_total.get("cached_input_tokens", 0),
|
usage_total.get("cached_input_tokens", 0),
|
||||||
usage_total.get("output_tokens", 0))
|
usage_total.get("output_tokens", 0))
|
||||||
used += micros or 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:
|
if used >= quota:
|
||||||
return ("quota_exceeded: 本月配额已用完(已用 %d / "
|
return ("quota_exceeded: 本月配额已用完(已用 %d / "
|
||||||
"配额 %d 微美元,%s)" % (used, quota, month))
|
"配额 %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
|
return None
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -205,6 +205,23 @@ def default_codex_bin():
|
|||||||
# 实时链路:单条通知里 item 的上限行数(防失控回合把文件写爆)。
|
# 实时链路:单条通知里 item 的上限行数(防失控回合把文件写爆)。
|
||||||
TRACE_MAX_LINES = 5000
|
TRACE_MAX_LINES = 5000
|
||||||
|
|
||||||
|
# agent shell 绝不能看到的变量(codex shell_environment_policy.exclude)。
|
||||||
|
# 实弹证实:不设此策略时 agent shell 直接继承 server 进程 env,
|
||||||
|
# XCM_API_KEY 在 local 路可被模型拿到;server 自己的中转调用走
|
||||||
|
# 进程 env(env_key),不受 shell 策略影响,所以只排 shell 不影响
|
||||||
|
# 模型调用。inherit 保持默认 all(PATH、代理变量等工作流变量不
|
||||||
|
# 动);gh 用文件登录不用 env,故不排 *_TOKEN;四个精确名是当前
|
||||||
|
# worker env 里的全部密钥类变量,三条模式兜底将来新增的同类变量。
|
||||||
|
SHELL_ENV_EXCLUDE = [
|
||||||
|
"XCM_API_KEY",
|
||||||
|
"CLOUD_DB_PASSWORD",
|
||||||
|
"GITHUB_OAUTH_CLIENT_SECRET",
|
||||||
|
"GITHUB_WEBHOOK_SECRET",
|
||||||
|
"*_KEY",
|
||||||
|
"*_SECRET",
|
||||||
|
"*_PASSWORD",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
def _summarize_trace_item(method, params):
|
def _summarize_trace_item(method, params):
|
||||||
"""app-server 通知 → 链路条目列表(实时链路用)。
|
"""app-server 通知 → 链路条目列表(实时链路用)。
|
||||||
@@ -302,6 +319,20 @@ def _append_trace_marker(state_path, task_name, entry):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def _server_command(codex_bin):
|
||||||
|
"""app-server 启动命令:-c 必须放在子命令之前。
|
||||||
|
|
||||||
|
shell_environment_policy.exclude 只过滤 agent shell 看到的
|
||||||
|
变量,不影响 server 进程自己的中转调用(env_key 直读进程 env)。
|
||||||
|
"""
|
||||||
|
return [
|
||||||
|
codex_bin,
|
||||||
|
"-c",
|
||||||
|
"shell_environment_policy.exclude=" + json.dumps(SHELL_ENV_EXCLUDE),
|
||||||
|
"app-server",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
class Server:
|
class Server:
|
||||||
"""单个 app-server 进程,多 thread 并发。"""
|
"""单个 app-server 进程,多 thread 并发。"""
|
||||||
|
|
||||||
@@ -365,7 +396,7 @@ class Server:
|
|||||||
os.makedirs(self.codex_home, exist_ok=True)
|
os.makedirs(self.codex_home, exist_ok=True)
|
||||||
env["CODEX_HOME"] = self.codex_home
|
env["CODEX_HOME"] = self.codex_home
|
||||||
self.proc = subprocess.Popen(
|
self.proc = subprocess.Popen(
|
||||||
[self.codex_bin, "app-server"],
|
_server_command(self.codex_bin),
|
||||||
stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
||||||
stderr=subprocess.DEVNULL, env=env,
|
stderr=subprocess.DEVNULL, env=env,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -52,6 +52,21 @@ class HostPureFunctionsTest(unittest.TestCase):
|
|||||||
missing = os.path.join(tmp, "missing.json")
|
missing = os.path.join(tmp, "missing.json")
|
||||||
self.assertEqual(host._load_config_file(missing), {})
|
self.assertEqual(host._load_config_file(missing), {})
|
||||||
|
|
||||||
|
def test_server_command_excludes_secrets_from_agent_shell(self):
|
||||||
|
cmd = host._server_command("/usr/bin/codex")
|
||||||
|
self.assertEqual(cmd[0], "/usr/bin/codex")
|
||||||
|
self.assertEqual(cmd[-1], "app-server")
|
||||||
|
self.assertIn("-c", cmd)
|
||||||
|
flag = cmd[cmd.index("-c") + 1]
|
||||||
|
self.assertTrue(
|
||||||
|
flag.startswith("shell_environment_policy.exclude="), flag)
|
||||||
|
excluded = json.loads(
|
||||||
|
flag.split("=", 1)[1])
|
||||||
|
for name in ("XCM_API_KEY", "CLOUD_DB_PASSWORD",
|
||||||
|
"GITHUB_OAUTH_CLIENT_SECRET", "GITHUB_WEBHOOK_SECRET",
|
||||||
|
"*_KEY", "*_SECRET", "*_PASSWORD"):
|
||||||
|
self.assertIn(name, excluded)
|
||||||
|
|
||||||
|
|
||||||
def _init_git_repo(tmp):
|
def _init_git_repo(tmp):
|
||||||
subprocess.run(["git", "init", "-q"], cwd=tmp, check=True)
|
subprocess.run(["git", "init", "-q"], cwd=tmp, check=True)
|
||||||
|
|||||||
Reference in new issue
Block a user