fix: Kimi collector targets correct membership/subscription tab, parses 3 usage metrics with dates

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,