diff --git a/usage-dashboard/usage-dashboard.py b/usage-dashboard/usage-dashboard.py index d9c4ac0..858d696 100644 --- a/usage-dashboard/usage-dashboard.py +++ b/usage-dashboard/usage-dashboard.py @@ -327,7 +327,9 @@ def latest_quota_age(): 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): # Dùng lock để tránh nhiều request đồng thời @@ -353,64 +355,118 @@ def refresh_quota(force=False): inserted = 0 with db_connect() as conn: for path in files: + basename = os.path.basename(path) try: auth = json.load(open(path)) token = auth.get("access_token") - email = auth.get("email") or os.path.basename(path) + email = auth.get("email") or basename if not token: - # ✅ LOG 2: Skip no token print(f"refresh_quota: skipping {email} (no access_token)", flush=True) continue - # ✅ LOG 3: Đang fetch quota cho email nào - print(f"refresh_quota: fetching quota for {email}...", flush=True) - req = urllib.request.Request( - "https://chatgpt.com/backend-api/wham/usage", - headers={ - "Authorization": "Bearer " + token, - "Accept": "application/json", - "User-Agent": "codex-cli", - }, - ) - with urllib.request.urlopen(req, timeout=20) as resp: - data = json.load(resp) - rl = data.get("rate_limit") or {} - primary = rl.get("primary_window") or {} - secondary = rl.get("secondary_window") or {} - primary_used = int(primary.get("used_percent") or 0) - secondary_used = int(secondary.get("used_percent") or 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"), - 1 if rl.get("allowed") else 0, - 1 if rl.get("limit_reached") else 0, - primary_used, - max(0, 100 - primary_used), - epoch_to_local(primary.get("reset_at")), - secondary_used, - max(0, 100 - secondary_used), - epoch_to_local(secondary.get("reset_at")), - str((data.get("credits") or {}).get("balance", "")), - json.dumps(data, ensure_ascii=False), - ), - ) - 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: 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( + "https://chatgpt.com/backend-api/wham/usage", + headers={ + "Authorization": "Bearer " + token, + "Accept": "application/json", + "User-Agent": "codex-cli", + }, + ) + with urllib.request.urlopen(req, timeout=20) as resp: + data = json.load(resp) + rl = data.get("rate_limit") or {} + primary = rl.get("primary_window") or {} + secondary = rl.get("secondary_window") or {} + primary_used = int(primary.get("used_percent") or 0) + secondary_used = int(secondary.get("used_percent") or 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"), + 1 if rl.get("allowed") else 0, + 1 if rl.get("limit_reached") else 0, + primary_used, + max(0, 100 - primary_used), + epoch_to_local(primary.get("reset_at")), + secondary_used, + max(0, 100 - secondary_used), + epoch_to_local(secondary.get("reset_at")), + str((data.get("credits") or {}).get("balance", "")), + json.dumps(data, ensure_ascii=False), + ), + ) + inserted += 1 + 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: - # ✅ 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) - # ✅ LOG 6: Tổng kết số quota snapshots đã insert + print(f"refresh_quota: inserted {inserted} quota snapshots", flush=True) return inserted finally: