add xai quota check

This commit is contained in:
2026-08-30 15:30:53 +07:00
parent 5517911783
commit e5b4dfdb07
+63 -7
View File
@@ -327,7 +327,9 @@ def latest_quota_age():
def auth_files(): def auth_files():
return sorted(glob.glob(os.path.join(JSON_DIR, "codex-*.json"))) codex_files = glob.glob(os.path.join(JSON_DIR, "codex-*.json"))
xai_files = glob.glob(os.path.join(JSON_DIR, "xai-*.json"))
return sorted(codex_files + xai_files)
def refresh_quota(force=False): def refresh_quota(force=False):
# Dùng lock để tránh nhiều request đồng thời # Dùng lock để tránh nhiều request đồng thời
@@ -353,16 +355,20 @@ def refresh_quota(force=False):
inserted = 0 inserted = 0
with db_connect() as conn: with db_connect() as conn:
for path in files: for path in files:
basename = os.path.basename(path)
try: try:
auth = json.load(open(path)) auth = json.load(open(path))
token = auth.get("access_token") token = auth.get("access_token")
email = auth.get("email") or os.path.basename(path) email = auth.get("email") or basename
if not token: if not token:
# ✅ LOG 2: Skip no token
print(f"refresh_quota: skipping {email} (no access_token)", flush=True) print(f"refresh_quota: skipping {email} (no access_token)", flush=True)
continue continue
# ✅ LOG 3: Đang fetch quota cho email nào
print(f"refresh_quota: fetching quota for {email}...", flush=True) print(f"refresh_quota: fetching quota for {email}...", flush=True)
# Kiểm tra tiền tố file để xác định endpoint
if basename.startswith("codex"):
# ChatGPT/Codex endpoint
req = urllib.request.Request( req = urllib.request.Request(
"https://chatgpt.com/backend-api/wham/usage", "https://chatgpt.com/backend-api/wham/usage",
headers={ headers={
@@ -405,12 +411,62 @@ def refresh_quota(force=False):
), ),
) )
inserted += 1 inserted += 1
# ✅ LOG 4: Lưu thành công, hiển thị plan type
print(f"refresh_quota: saved quota for {email} (plan: {data.get('plan_type')})", flush=True) print(f"refresh_quota: saved quota for {email} (plan: {data.get('plan_type')})", flush=True)
elif basename.startswith("xai"):
# xAI endpoint
req = urllib.request.Request(
"https://api.x.ai/v1/api-key",
headers={
"Authorization": "Bearer " + token,
"Accept": "application/json",
"User-Agent": "codex-cli",
},
)
with urllib.request.urlopen(req, timeout=20) as resp:
data = json.load(resp)
# xAI trả về format khác, cần parse theo cấu trúc của nó
# Giả sử xAI trả về: {"rate_limit": {"remaining": 100, "limit": 1000, "reset_at": ...}}
rl = data.get("rate_limit") or {}
remaining = int(rl.get("remaining") or 0)
limit = int(rl.get("limit") or 100)
used_percent = int((limit - remaining) * 100 / limit) if limit > 0 else 0
conn.execute(
"""
INSERT INTO quota_snapshots (
timestamp,ts_epoch,email,plan,allowed,limit_reached,
primary_used_percent,primary_remaining_percent,primary_reset_at,
secondary_used_percent,secondary_remaining_percent,secondary_reset_at,
credits_balance,raw_json
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)
""",
(
now.isoformat(),
now.timestamp(),
email,
data.get("plan_type", "xai"),
1 if not rl.get("limit_reached") else 0,
1 if rl.get("limit_reached") else 0,
used_percent,
max(0, 100 - used_percent),
epoch_to_local(rl.get("reset_at")),
0, # xAI không có secondary window
0,
"",
str(data.get("credits_balance", "")),
json.dumps(data, ensure_ascii=False),
),
)
inserted += 1
print(f"refresh_quota: saved quota for {email} (provider: xAI)", flush=True)
else:
print(f"refresh_quota: skipping {email} (unknown prefix: {basename})", flush=True)
except (OSError, urllib.error.URLError, urllib.error.HTTPError, json.JSONDecodeError, KeyError) as exc: except (OSError, urllib.error.URLError, urllib.error.HTTPError, json.JSONDecodeError, KeyError) as exc:
# ✅ LOG 5: Lỗi khi fetch (đã có sẵn, thêm flush=True)
print(f"quota refresh failed for {path}: {exc}", file=sys.stderr, flush=True) print(f"quota refresh failed for {path}: {exc}", file=sys.stderr, flush=True)
# ✅ LOG 6: Tổng kết số quota snapshots đã insert
print(f"refresh_quota: inserted {inserted} quota snapshots", flush=True) print(f"refresh_quota: inserted {inserted} quota snapshots", flush=True)
return inserted return inserted
finally: finally: