games-development-ai/game-runtime/tools/newapi_pool_provision.py
lili a54ba9bb42 docs(neice): 内测闭环三份设计档(WU1注册登录/WU2 newapi额度预置池/WU3 dify节点) + ops 池预置脚本
阶段0 三份 feature-design-doc 过对抗双评审+主控终审;WU2 按 S0 实测改预置池机制。
ops 脚本离线预建 new-api user+token+¥100 灌 game-cloud 池表。

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-07 18:05:00 -07:00

422 lines
21 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""newapi_pool_provision.py -- 内测 new-api per-user ¥100 额度池离线预置脚本WU2 · 接线点 B
权威设计docs/agent-specs/2026-07-07-内测-WU2-newapi额度接入-设计.md §3.3(离线 ops 预置脚本)
+ §3.5(池表 newapi_quota_pool 契约)。
它做什么
--------
在 mini-infra 上一次性预建 N 个专属 (new-api user + token + ¥100 额度) 账户,
把每个条目 emit 成可导入 game-cloud 池表 newapi_quota_pool 的 JSON + SQL供 WU2 落地后灌数。
provision 全部离线做S0 实测坐实new-api admin 令牌不能替他人建 token、access_token
不可经 API 设per-user token 必须直连其 postgres运行时零 new-api admin 耦合。
每个池条目四步建成S0 坐实的唯一可行路径)
------------------------------------------------
1. POST /api/user/root 令牌建用户username = 确定性 `<前缀><序号>`;从 postgres 查 uid。
2. postgres 直写 UPDATE users SET access_token=<确定性 32 位串> WHERE id=uid
—— access_token 无法经 API 设,又是下一步「以该用户身份建 token」的前提。
3. POST /api/token/(以该用户 access_token + New-Api-User:uid 建 token
remain_quota=¥100 折算、unlimited_quota=false、expired_time=-1不过期从 postgres 查 token_id/key。
4. PUT /api/user/root 令牌)设 user.quota=¥100 折算(账户自洽;权威余额闸仍是 token.remain_quota
幂等(可重跑补池)
------------------
按确定性 username 去重:重跑时若用户已存在则复用其 uid逐步「补齐」缺失的 access_token /
token / quotaensure 语义),不重复建号。补池 = 提高 --count 重跑。单条失败只记该条、不污染后续。
外部交互红线
------------
- HTTP连接/读超时urlopen timeout=读10s内部连粒度由 socket 兜),非 2xx / success=false 归错,
5xx/超时重试 2 次指数退避4xx 不重试全程绕系统代理ProxyHandler({}),内网直连 100.64.x
- postgres经 `docker exec infra-postgres psql` 执行S0 验证路径),失败即抛、该条标记失败。
- 日志token / access_token 一律脱敏(前后各留 4 位);可追溯(时间戳 + 步骤 + 条目序号)。
凭据(不硬编码)
----------------
- postgres 密码:环境变量 NEWAPI_PG_PASSWORD必填见 docs/内网凭据与端点.md
- root 管理令牌:环境变量 NEWAPI_ROOT_TOKEN选填未设则自动从 postgres users.id=1 读。
用法(在 mini-infra 上跑docker 本地可用)
------------------------------------------
export NEWAPI_PG_PASSWORD=<见凭据文档>
# 预置 2 条测试条目验证用前缀明确可辨dry-run 先看:
python3 newapi_pool_provision.py --count 2 --prefix neice_s0test_ --dry-run
python3 newapi_pool_provision.py --count 2 --prefix neice_s0test_
# 清理测试条目(删 tokens + users 行):
python3 newapi_pool_provision.py --cleanup 'neice_s0test_%'
"""
from __future__ import annotations
import argparse
import hashlib
import json
import logging
import os
import re
import secrets
import subprocess
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path
# ─── 常量 ────────────────────────────────────────────────────────────────
NEWAPI_BASE = os.environ.get("NEWAPI_BASE", "http://localhost:3000") # 脚本宿主 mini-infra本地直连
PG_CONTAINER = os.environ.get("NEWAPI_PG_CONTAINER", "infra-postgres")
PG_DB = os.environ.get("NEWAPI_PG_DB", "new-api")
PG_USER = os.environ.get("NEWAPI_PG_USER", "root")
ROOT_UID = 1 # new-api root 用户 idqingse持系统管理 access_token
HTTP_CONNECT_TIMEOUT = 5.0
HTTP_READ_TIMEOUT = 10.0
HTTP_RETRIES = 2 # 5xx / 超时最多重试 2 次
USERNAME_RE = re.compile(r"^[A-Za-z0-9_]+$") # 只允许安全 username杜绝 SQL/命令注入
POOL_TABLE = "newapi_quota_pool"
log = logging.getLogger("newapi_pool")
def _mask(token: str | None) -> str:
"""脱敏敏感串:前后各留 4 位中间打码。用于日志红线token/access_token 不整条打日志)。"""
if not token:
return "<empty>"
if len(token) <= 10:
return token[:2] + "***"
return f"{token[:4]}{token[-4:]}(len={len(token)})"
# ─── postgres 直连(经 docker exec psqlS0 验证路径) ──────────────────────
class Postgres:
"""封装对 new-api postgres 的只读/写操作。所有写走事务语义(单条 SQL 天然原子)。"""
def __init__(self, password: str):
if not password:
raise RuntimeError(
"缺少 NEWAPI_PG_PASSWORD 环境变量——无法连接 new-api postgres。"
"见 docs/内网凭据与端点.md new-api 段。"
)
self._password = password
def _run(self, sql: str) -> str:
"""执行一条 SQL返回 tuples-only、unaligned、以 | 分隔的原始 stdout。失败即抛。"""
cmd = [
"docker", "exec",
"-e", f"PGPASSWORD={self._password}",
PG_CONTAINER,
"psql", "-U", PG_USER, "-d", PG_DB,
"-t", "-A", "-F", "|", "-c", sql,
]
try:
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
except subprocess.TimeoutExpired as e:
raise RuntimeError(f"postgres 执行超时:{sql[:80]}") from e
if proc.returncode != 0:
# 可追溯错误日志:暴露 psql stderr但不含敏感值SQL 里的 access_token 已在调用侧脱敏日志)
raise RuntimeError(f"postgres 执行失败 rc={proc.returncode}: {proc.stderr.strip()}")
return proc.stdout.strip()
def query_root_token(self) -> str:
out = self._run(f"SELECT access_token FROM users WHERE id={ROOT_UID};")
if not out:
raise RuntimeError("postgres 未查到 root(id=1) access_token")
return out.splitlines()[0].strip()
def query_user_id(self, username: str) -> int | None:
out = self._run(f"SELECT id FROM users WHERE username='{username}';")
line = out.splitlines()[0].strip() if out else ""
return int(line) if line else None
def query_access_token(self, uid: int) -> str:
out = self._run(f"SELECT access_token FROM users WHERE id={uid};")
return out.splitlines()[0].strip() if out else ""
def query_user_quota(self, uid: int) -> int | None:
out = self._run(f"SELECT quota FROM users WHERE id={uid};")
line = out.splitlines()[0].strip() if out else ""
return int(line) if line else None
def set_access_token(self, uid: int, access_token: str) -> None:
# access_token 无法经 API 设,必须 DB 直写S0 坐实的关键步)
self._run(f"UPDATE users SET access_token='{access_token}' WHERE id={uid};")
def query_token(self, uid: int, name: str) -> tuple[int, str] | None:
"""按 user_id + token name 查 token返回 (token_id, token_key)。"""
out = self._run(
f"SELECT id, key FROM tokens WHERE user_id={uid} AND name='{name}' ORDER BY id LIMIT 1;"
)
line = out.splitlines()[0].strip() if out else ""
if not line or "|" not in line:
return None
tid, key = line.split("|", 1)
return int(tid), key.strip()
def delete_user_cascade(self, username_like: str) -> tuple[int, int]:
"""清理:按 username LIKE 删 users + 其 tokens。返回 (删 tokens 数, 删 users 数)。
安全护栏:拒绝任何可能命中 root(id<=1) 的模式;只允许显式 % 通配。"""
# 先查将被删的 uid 列表,护栏校验
uids_out = self._run(f"SELECT id FROM users WHERE username LIKE '{username_like}';")
uids = [int(x) for x in uids_out.splitlines() if x.strip()]
if not uids:
return (0, 0)
if any(u <= ROOT_UID for u in uids):
raise RuntimeError(f"清理护栏:模式 {username_like} 命中受保护用户 id<={ROOT_UID},拒绝执行")
uid_csv = ",".join(str(u) for u in uids)
tok_out = self._run(f"WITH d AS (DELETE FROM tokens WHERE user_id IN ({uid_csv}) RETURNING 1) SELECT count(*) FROM d;")
usr_out = self._run(f"WITH d AS (DELETE FROM users WHERE id IN ({uid_csv}) RETURNING 1) SELECT count(*) FROM d;")
return (int(tok_out.strip() or 0), int(usr_out.strip() or 0))
# ─── new-api HTTPstdlib urllib绕系统代理带超时+重试) ──────────────────
_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({})) # 空 ProxyHandler = 禁用系统代理
def _http_json(method: str, path: str, token: str, newapi_user: int, body: dict | None = None) -> dict:
"""对 new-api 发一个 JSON 请求。红线:绕代理 / 超时 / 非 2xx 归错 / 5xx 重试。
path 必须带尾斜杠new-api gin RedirectTrailingSlash否则 307
鉴权头Authorization: Bearer <token> + New-Api-User: <uid>(缺后者即 401
"""
url = NEWAPI_BASE.rstrip("/") + path
data = json.dumps(body).encode("utf-8") if body is not None else None
headers = {
"Authorization": f"Bearer {token}",
"New-Api-User": str(newapi_user),
"Content-Type": "application/json",
}
last_err: Exception | None = None
for attempt in range(HTTP_RETRIES + 1):
req = urllib.request.Request(url, data=data, headers=headers, method=method)
try:
with _OPENER.open(req, timeout=HTTP_READ_TIMEOUT) as resp:
raw = resp.read().decode("utf-8")
payload = json.loads(raw) if raw else {}
# new-api 标准返回体 {success,message,data}success=false 视为业务错
if isinstance(payload, dict) and payload.get("success") is False:
raise RuntimeError(f"{method} {path} 业务失败: {payload.get('message')}")
return payload
except urllib.error.HTTPError as e:
body_txt = e.read().decode("utf-8", "ignore")[:200]
last_err = RuntimeError(f"{method} {path} HTTP {e.code}: {body_txt}")
if 500 <= e.code < 600 and attempt < HTTP_RETRIES:
time.sleep(0.5 * (2 ** attempt))
log.warning("[http] %s %s 5xx 重试 %d/%d", method, path, attempt + 1, HTTP_RETRIES)
continue
raise last_err # 4xx 不重试
except (urllib.error.URLError, TimeoutError, OSError) as e:
last_err = RuntimeError(f"{method} {path} 网络错误: {e}")
if attempt < HTTP_RETRIES:
time.sleep(0.5 * (2 ** attempt))
log.warning("[http] %s %s 网络错重试 %d/%d: %s", method, path, attempt + 1, HTTP_RETRIES, e)
continue
raise last_err
raise last_err or RuntimeError("unreachable")
# ─── ¥ ↔ quota 折算(口径与 cost.py 同源,读网关而非硬编码) ─────────────────
def fetch_conversion() -> tuple[int, float]:
"""GET /api/status 拿 quota_per_unit / usd_exchange_rate无需鉴权。缺 usd 用 7.3 默认。"""
url = NEWAPI_BASE.rstrip("/") + "/api/status"
req = urllib.request.Request(url, method="GET")
with _OPENER.open(req, timeout=HTTP_READ_TIMEOUT) as resp:
data = json.loads(resp.read().decode("utf-8")).get("data", {})
qpu = int(data.get("quota_per_unit") or 500000)
usd = data.get("usd_exchange_rate")
usd = float(usd) if usd else 7.3 # DB options 未覆盖时默认 7.3
return qpu, usd
def compute_grant(yuan: float, qpu: int, usd: float) -> int:
"""¥ → quotaround(yuan / usd × qpu)。¥100 @ 500000/7.3 ≈ 6,849,315。"""
return round(yuan / usd * qpu)
# ─── 单条目预置ensure 语义,幂等可补齐) ──────────────────────────────────
def provision_one(pg: Postgres, root_token: str, username: str, grant: int,
qpu: int, usd: float, dry_run: bool) -> dict:
"""预置/补齐一个池条目,返回条目 dict对齐 §3.5 池表契约)。任一步失败即抛,由调用方记为失败。"""
if not USERNAME_RE.match(username):
raise RuntimeError(f"非法 username: {username}")
if dry_run:
log.info("[dry-run] 计划预置 %sgrant=%d, unlimited=false, expired=-1", username, grant)
return {"username": username, "dry_run": True, "grant_quota": grant}
# 步骤 1建用户幂等——已存在则复用
uid = pg.query_user_id(username)
if uid is None:
password = secrets.token_urlsafe(12) # 随机密码 16 位new-api Password max=20仅占位从不密码登录
_http_json("POST", "/api/user/", root_token, ROOT_UID, {
"username": username,
"password": password,
"display_name": username[:20], # new-api DisplayName max=20直接用 username截断兜底
})
uid = pg.query_user_id(username)
if uid is None:
raise RuntimeError(f"建用户 {username} 后 postgres 未查到 uid")
log.info("[1/4] 建用户 %s → uid=%d", username, uid)
else:
log.info("[1/4] 用户 %s 已存在 uid=%d(复用)", username, uid)
# 步骤 2DB 直写 access_token确定性可重复 UPDATE 幂等)
access_token = hashlib.sha256(f"neice-quota-pool::v1::{username}".encode()).hexdigest()[:32]
existing_at = pg.query_access_token(uid)
if existing_at != access_token:
pg.set_access_token(uid, access_token)
log.info("[2/4] DB 写 access_token uid=%d%s", uid, _mask(access_token))
else:
log.info("[2/4] access_token uid=%d 已就绪 %s(复用)", uid, _mask(access_token))
# 步骤 3以该用户身份建 token幂等——已存在同名 token 则复用)
tok = pg.query_token(uid, username)
if tok is None:
_http_json("POST", "/api/token/", access_token, uid, {
"name": username,
"remain_quota": grant,
"unlimited_quota": False,
"expired_time": -1, # 不过期,避免自伤式过期(设计 §3.9
})
tok = pg.query_token(uid, username)
if tok is None:
raise RuntimeError(f"建 token 后 postgres 未查到 user={uid} name={username}")
log.info("[3/4] 建 token uid=%d → token_id=%d key=%s", uid, tok[0], _mask(tok[1]))
else:
log.info("[3/4] token uid=%d name=%s 已存在 token_id=%d(复用)", uid, username, tok[0])
token_id, token_key = tok
# 步骤 4设 user.quotaroot 令牌;账户自洽,权威余额闸仍是 token.remain_quota
cur_quota = pg.query_user_quota(uid)
if cur_quota != grant:
_http_json("PUT", "/api/user/", root_token, ROOT_UID, {
"id": uid,
"username": username,
"display_name": username[:20],
"quota": grant,
"group": "default",
})
log.info("[4/4] 设 user.quota uid=%d%d", uid, grant)
else:
log.info("[4/4] user.quota uid=%d 已=%d(复用)", uid, grant)
# 组装池表条目§3.5 契约字段)
return {
"newapi_user_id": uid,
"newapi_token_id": token_id,
"newapi_token_key": token_key,
"grant_quota": grant,
"quota_per_unit_snapshot": qpu,
"usd_rate_snapshot": usd,
"status": "FREE",
"username": username, # 仅审计参考,非池表列
}
# ─── emitJSON + SQL对齐 §3.5,供 WU2 落地后导入 game-cloud ─────────────
def emit_outputs(entries: list[dict], out_dir: Path, prefix: str) -> tuple[Path, Path]:
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
out_dir.mkdir(parents=True, exist_ok=True)
stem = f"newapi_pool_{prefix.rstrip('_')}_{ts}"
json_path = out_dir / f"{stem}.json"
sql_path = out_dir / f"{stem}.sql"
# 只保留池表列(去掉 username 审计字段)
pool_cols = ["newapi_user_id", "newapi_token_id", "newapi_token_key",
"grant_quota", "quota_per_unit_snapshot", "usd_rate_snapshot", "status"]
json_rows = [{k: e[k] for k in pool_cols} for e in entries]
json_path.write_text(json.dumps(json_rows, ensure_ascii=False, indent=2), encoding="utf-8")
# SQLINSERT ... ON DUPLICATE KEY UPDATE去重键 uk(newapi_user_id)§3.3(a)
# 刻意不在 ON DUPLICATE 里改 status/claimed_*——重跑不把已 CLAIMED 条目刷回 FREE防免费续杯 / 防误 unclaim
lines = [
f"-- 内测 new-api 额度池预置产物 {ts},共 {len(json_rows)} 条。",
f"-- 对齐 docs/agent-specs/2026-07-07-内测-WU2-newapi额度接入-设计.md §3.5 池表 {POOL_TABLE}V32 待落)。",
f"-- 导入前提game-cloud 已建 {POOL_TABLE} 表(含 uk(newapi_user_id)、审计列默认值)。",
f"-- ON DUPLICATE 只更 token/额度快照,不动 status/claimed_*(幂等补池不 unclaim 已绑条目)。",
]
for r in json_rows:
key_sql = r["newapi_token_key"].replace("'", "''")
lines.append(
f"INSERT INTO {POOL_TABLE} "
f"(newapi_user_id, newapi_token_id, newapi_token_key, grant_quota, "
f"quota_per_unit_snapshot, usd_rate_snapshot, status) VALUES "
f"({r['newapi_user_id']}, {r['newapi_token_id']}, '{key_sql}', {r['grant_quota']}, "
f"{r['quota_per_unit_snapshot']}, {r['usd_rate_snapshot']}, '{r['status']}') "
f"ON DUPLICATE KEY UPDATE "
f"newapi_token_id=VALUES(newapi_token_id), newapi_token_key=VALUES(newapi_token_key), "
f"grant_quota=VALUES(grant_quota), quota_per_unit_snapshot=VALUES(quota_per_unit_snapshot), "
f"usd_rate_snapshot=VALUES(usd_rate_snapshot);"
)
sql_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
return json_path, sql_path
# ─── 主流程 ──────────────────────────────────────────────────────────────
def cmd_provision(args) -> int:
pg = Postgres(os.environ.get("NEWAPI_PG_PASSWORD", ""))
root_token = os.environ.get("NEWAPI_ROOT_TOKEN") or pg.query_root_token()
log.info("root 管理令牌就绪 %s", _mask(root_token))
qpu, usd = fetch_conversion()
grant = compute_grant(args.yuan, qpu, usd)
log.info("折算:¥%.0f @ quota_per_unit=%d usd=%.4f → grant_quota=%d", args.yuan, qpu, usd, grant)
entries: list[dict] = []
failures: list[dict] = []
for i in range(args.start, args.start + args.count):
username = f"{args.prefix}{i:03d}"
try:
entries.append(provision_one(pg, root_token, username, grant, qpu, usd, args.dry_run))
except Exception as e: # 单条失败不污染后续(红线)
log.error("[条目 %s] 预置失败:%s", username, e)
failures.append({"username": username, "error": str(e)})
log.info("完成:成功 %d / 失败 %d", len(entries), len(failures))
if args.dry_run:
log.info("dry-run 结束,未写 new-api、未 emit 文件。")
return 0 if not failures else 1
if entries:
json_path, sql_path = emit_outputs(entries, Path(args.out_dir), args.prefix)
log.info("emit JSON → %s", json_path)
log.info("emit SQL → %s", sql_path)
if failures:
log.warning("失败条目:%s", json.dumps(failures, ensure_ascii=False))
return 1
return 0
def cmd_cleanup(args) -> int:
"""清理:按 username LIKE 模式删 tokens + users护栏拒绝命中 root"""
pg = Postgres(os.environ.get("NEWAPI_PG_PASSWORD", ""))
ntok, nusr = pg.delete_user_cascade(args.cleanup)
log.info("清理模式 %s:删 tokens %d 行 / users %d", args.cleanup, ntok, nusr)
return 0
def main() -> int:
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s %(message)s",
stream=sys.stderr,
)
p = argparse.ArgumentParser(description="内测 new-api per-user ¥100 额度池离线预置")
p.add_argument("--count", type=int, default=1, help="目标条目数 N本次预置到 start..start+count")
p.add_argument("--start", type=int, default=1, help="起始序号(默认 1")
p.add_argument("--prefix", default="neice_", help="username 前缀(默认 neice_测试用 neice_s0test_")
p.add_argument("--yuan", type=float, default=100.0, help="每条目额度(元,默认 100")
p.add_argument("--out-dir", default="./newapi-pool-out", help="JSON/SQL 产物目录")
p.add_argument("--dry-run", action="store_true", help="只打印计划,不写 new-api、不 emit")
p.add_argument("--cleanup", metavar="LIKE_PATTERN",
help="清理模式:删匹配 username LIKE 的 users+tokens'neice_s0test_%%'")
args = p.parse_args()
if args.cleanup:
return cmd_cleanup(args)
return cmd_provision(args)
if __name__ == "__main__":
sys.exit(main())