From 84423565f8a1eb00431c132e9ed58e30a6d6c11d Mon Sep 17 00:00:00 2001 From: hmo Date: Sun, 19 Jul 2026 21:26:40 +0800 Subject: [PATCH] fix: Kimi collector targets correct membership/subscription tab, parses 3 usage metrics with dates --- gateway/scripts/usage_collector_kimi.py | 111 +++++++++++++++++------- 1 file changed, 79 insertions(+), 32 deletions(-) diff --git a/gateway/scripts/usage_collector_kimi.py b/gateway/scripts/usage_collector_kimi.py index cafa67c..f63f6e5 100644 --- a/gateway/scripts/usage_collector_kimi.py +++ b/gateway/scripts/usage_collector_kimi.py @@ -52,33 +52,33 @@ def cdp_eval(target_id, js, timeout=10): def main(): print("[kimi] Starting collection...") - # 1. Find or create a Kimi tab + # 1. Find tab with the membership subscription page targets = cdp_get("/targets") or [] kimi_tab = None for t in targets: url = t.get("url", "") - if "kimi.com" in url or "moonshot.cn" in url: + if "membership/subscription" in url or "kimi.com/membership" in url: kimi_tab = t break if kimi_tab: target_id = kimi_tab.get("targetId") or kimi_tab.get("id") - print(f"[kimi] Found existing tab: {target_id}") - # Navigate to quota page if needed + print(f"[kimi] Found quota tab: {target_id}") + # Navigate if needed cur_url = kimi_tab.get("url", "") - if "membership" not in cur_url and "subscription" not in cur_url: - print(f"[kimi] Navigating to quota page...") + if "tab=quota" not in cur_url: + print(f"[kimi] Navigating to quota tab...") cdp_get(f"/navigate?target={target_id}&url={KIMI_CONSOLE}") - time.sleep(3) + time.sleep(5) else: - print(f"[kimi] Opening new tab...") + print(f"[kimi] No quota tab found, opening new tab...") result = cdp_get(f"/new?url={KIMI_CONSOLE}") if isinstance(result, dict): target_id = result.get("targetId") or result.get("id", "") else: target_id = str(result) print(f"[kimi] New tab: {target_id}") - time.sleep(4) + time.sleep(5) if not target_id: print("[kimi] ERROR: no target ID") @@ -101,19 +101,66 @@ def main(): # Parse the text in Python import re - # Find all "N%" followed by "N 小时后重置" patterns - # The page has two: first = 本周用量, second = 频率明细 - pct_pattern = re.findall(r'(\d+)%', text) - hour_pattern = re.findall(r'(\d+)\s*小时后重置', text) + # Membership page has three metrics: + # 总用量 (monthly): N%, reset YYYY-MM-DD + # 5小时用量 (5h): Code N%, reset MM-DD HH:MM + # 7天用量 (7d): Code N%, reset MM-DD HH:MM + # The percentages appear in order: total, 5h, 7d - weekly_pct = int(pct_pattern[0]) if len(pct_pattern) > 0 else 0 - weekly_hours = int(hour_pattern[0]) if len(hour_pattern) > 0 else 0 - rate_pct = int(pct_pattern[1]) if len(pct_pattern) > 1 else 0 - rate_hours = int(hour_pattern[1]) if len(hour_pattern) > 1 else 0 + # Extract all percentage values with context + pct_matches = list(re.finditer(r'(\d+(?:\.\d+)?)%', text)) + # Extract dates in format MM-DD HH:MM or YYYY-MM-DD + date_matches = list(re.finditer(r'(\d{4}-\d{2}-\d{2})|(\d{2}-\d{2}\s+\d{2}:\d{2})', text)) - print(f"[kimi] Found {len(pct_pattern)} pct values, {len(hour_pattern)} hour values") + # Build list of (pct, date) pairs based on position in text + entries = [] + for m in pct_matches: + pct_val = float(m.group(1)) + pct_pos = m.start() + # Find nearest date after this percentage + best_date = None + for d in date_matches: + if d.start() > pct_pos: + best_date = d.group(1) or d.group(2) + break + entries.append((pct_val, best_date, pct_pos)) - print(f"[kimi] Weekly: {weekly_pct}% / {weekly_hours}h reset | Rate: {rate_pct}% / {rate_hours}h reset") + # Assign by order: first = total (monthly), second = 5h, third = 7d + def parse_date(dstr): + """Parse date string to seconds remaining.""" + if not dstr: + return 0 + now = datetime.now(timezone.utc) + try: + if '-' in dstr and ':' not in dstr: + # YYYY-MM-DD + d = datetime.strptime(dstr, "%Y-%m-%d") + return max(0, int((d.replace(tzinfo=timezone.utc) - now).total_seconds())) + elif '-' in dstr and ':' in dstr: + # MM-DD HH:MM — assume current year + year = now.year + d = datetime.strptime(dstr, "%m-%d %H:%M").replace(year=year, tzinfo=timezone.utc) + if d < now: + d = d.replace(year=year + 1) + return max(0, int((d - now).total_seconds())) + except Exception: + pass + return 0 + + def get_status(pct): + if pct >= 95: return "rate-limited" + if pct >= 80: return "warn" + return "ok" + + # Build metrics + total_pct = entries[0][0] if len(entries) > 0 else 0 + total_reset = parse_date(entries[0][1]) if len(entries) > 0 else 0 + five_hr_pct = entries[1][0] if len(entries) > 1 else 0 + five_hr_reset = parse_date(entries[1][1]) if len(entries) > 1 else 0 + seven_d_pct = entries[2][0] if len(entries) > 2 else 0 + seven_d_reset = parse_date(entries[2][1]) if len(entries) > 2 else 0 + + print(f"[kimi] Monthly: {total_pct}% | 5h: {five_hr_pct}% | 7d: {seven_d_pct}%") # Also take a screenshot for debugging try: @@ -122,12 +169,8 @@ def main(): except: pass - # Build output with Python-parsed data - wp = weekly_pct or 0 - wh = weekly_hours or 0 - rp = rate_pct or 0 - rh = rate_hours or 0 - + # Build output with correctly mapped metrics + # rolling → 5小时用量, weekly → 7天用量, monthly → 总用量 output = { "ok": True, "last_refresh_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"), @@ -138,16 +181,20 @@ def main(): "provider": PROVIDER, "workspace_id": "kimi", "rolling": { - "usage_percent": rp, - "reset_in_sec": rh * 3600, - "status": "ok" if rp < 80 else ("warn" if rp < 95 else "rate-limited"), + "usage_percent": five_hr_pct, + "reset_in_sec": five_hr_reset, + "status": get_status(five_hr_pct), }, "weekly": { - "usage_percent": wp, - "reset_in_sec": wh * 3600, - "status": "ok" if wp < 80 else ("warn" if wp < 95 else "rate-limited"), + "usage_percent": seven_d_pct, + "reset_in_sec": seven_d_reset, + "status": get_status(seven_d_pct), + }, + "monthly": { + "usage_percent": total_pct, + "reset_in_sec": total_reset, + "status": get_status(total_pct), }, - "monthly": None, "session_expired": False, "last_update_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"), "error": None,