update usage-dashboard theo flow moi cua cliproxyapi

usage-dashboard/usage-dashboard.py

@@ -36,6 +36,8 @@ DEFAULT_CONFIG = {
     "quota_refresh_seconds": 300,
     "dashboard_host": "0.0.0.0",
     "dashboard_port": 8320,
+    "queue_channel": "usage",   # CLIProxy dùng pub/sub trên channel này
+    "collector_mode": "subscribe",  # "subscribe" (pub/sub) hoặc "rpop" (queue)
 }
 
 
@@ -193,10 +195,39 @@ class RespClient:
             return [self.read() for _ in range(count)]
         raise RuntimeError(f"Unknown RESP prefix: {line!r}")
 
-    def rpop(self, count=100):
+    def cmd(self, *parts):
+        """Gửi bất kỳ command nào và trả về response (không raise lỗi)."""
+        self.send(*parts)
+        try:
+            return self.read()
+        except RuntimeError as e:
+            return f"ERR: {e}"
+
+    def subscribe(self, channel="usage"):
+        """Gửi SUBSCRIBE và đọc reply xác nhận."""
+        self.send("SUBSCRIBE", channel)
+        reply = self.read()  # ['subscribe', channel, 1]
+        return reply
+
+    def read_message(self):
+        """
+        Đọc 1 pub/sub message. Trả về data string nếu là message thật,
+        None nếu là subscribe confirmation hay message hệ thống.
+        Format RESP: ['message', channel, data]
+        """
+        msg = self.read()
+        if not isinstance(msg, list) or len(msg) < 3:
+            return None
+        kind = msg[0] if isinstance(msg[0], str) else ""
+        if kind == "message":
+            return msg[2]  # data payload
+        return None  # 'subscribe'/'unsubscribe'/ping
+
+    def rpop(self, count=100, channel="usage"):
+        """Fallback RPOP — dùng khi server hỗ trợ queue mode."""
         result = []
         for _ in range(count):
-            self.send("RPOP", "queue")
+            self.send("RPOP", channel)
             item = self.read()
             if item is None:
                 break
@@ -374,8 +405,10 @@ def epoch_to_local(value):
 
 
 def debug_collect():
-    """Test one poll cycle and print raw queue items — does NOT write to DB."""
+    """Test pub/sub subscribe, nhận tối đa 5 event rồi thoát — không ghi DB."""
+    import select as _select
     cfg = load_config()
+    channel = cfg.get("queue_channel", "usage")
     print(f"[debug] Config: {cfg}", flush=True)
     print(f"[debug] Connecting to {cfg['cliproxy_host']}:{cfg['cliproxy_port']} ...", flush=True)
     try:
@@ -385,45 +418,68 @@ def debug_collect():
         return
     print("[debug] AUTH OK", flush=True)
     try:
-        raw_items = client.rpop(10)
-        print(f"[debug] rpop returned {len(raw_items)} item(s)", flush=True)
-        for i, raw in enumerate(raw_items):
-            print(f"[debug] item[{i}]: {raw[:300]}", flush=True)
+        reply = client.subscribe(channel)
+        print(f"[debug] SUBSCRIBE {channel!r} → {reply}", flush=True)
+        print(f"[debug] Chờ event từ CLIProxy (timeout 15s, tối đa 5 event)...", flush=True)
+        client.sock.settimeout(None)  # blocking, dùng select
+
+        received = 0
+        deadline = time.time() + 15
+        while received < 5:
+            remaining = deadline - time.time()
+            if remaining <= 0:
+                print("[debug] Timeout — không có event nào trong 15s.", flush=True)
+                print("[debug] Hãy thực hiện 1 request qua CLIProxy rồi chạy lại.", flush=True)
+                break
+            readable, _, _ = _select.select([client.sock], [], [], remaining)
+            if not readable:
+                print("[debug] Timeout — không có event nào trong 15s.", flush=True)
+                print("[debug] Hãy thực hiện 1 request qua CLIProxy rồi chạy lại.", flush=True)
+                break
+            data = client.read_message()
+            if data is None:
+                continue  # confirmation frame
+            # Bỏ qua heartbeat
             try:
-                payload = json.loads(raw)
-                print(f"[debug] parsed OK — keys: {list(payload.keys())}", flush=True)
+                obj = json.loads(data)
+                if isinstance(obj, dict) and set(obj.keys()) == {"support_refresh"}:
+                    print(f"[debug] (heartbeat nhận được, đang chờ event thật...)", flush=True)
+                    continue
+            except Exception:
+                pass
+            received += 1
+            print(f"[debug] event[{received}] raw: {data[:300]}", flush=True)
+            try:
+                payload = json.loads(data)
+                print(f"[debug]   keys: {list(payload.keys())}", flush=True)
                 tokens = payload.get("tokens") or {}
-                print(f"[debug]   tokens: {tokens}", flush=True)
-                print(f"[debug]   model: {payload.get('model')}, source: {payload.get('source')}", flush=True)
+                print(f"[debug]   model={payload.get('model')} source={payload.get('source')} tokens={tokens}", flush=True)
             except Exception as e:
-                print(f"[debug] parse error: {e}", file=sys.stderr, flush=True)
+                print(f"[debug]   parse error: {e}", file=sys.stderr, flush=True)
+        if received == 0:
+            print("[debug] Không nhận được event nào.", flush=True)
+        else:
+            print(f"[debug] Nhận được {received} event(s) thành công ✅", flush=True)
     finally:
         client.close()
-    if not raw_items:
-        print("[debug] Queue is EMPTY — no events to collect. Make sure CLIProxy is running and receiving requests.", flush=True)
 
 
 def collect_forever():
     init_db()
     cfg = load_config()
-    last_quota = 0
-    print(f"collector: connecting to {cfg['cliproxy_host']}:{cfg['cliproxy_port']}", flush=True)
+    channel = cfg.get("queue_channel", "usage")
+    mode = cfg.get("collector_mode", "subscribe")
+    print(f"collector: mode={mode!r} channel={channel!r} → {cfg['cliproxy_host']}:{cfg['cliproxy_port']}", flush=True)
+
     while True:
         try:
             client = RespClient(cfg["cliproxy_host"], cfg["cliproxy_port"], cfg["management_key"])
             print("collector: connected and authenticated OK", flush=True)
             try:
-                while True:
-                    raw_items = client.rpop(100)
-                    if raw_items:
-                        print(f"collector: got {len(raw_items)} raw item(s) from queue", flush=True)
-                        inserted = insert_usage(raw_items)
-                        print(f"collector: inserted {inserted}/{len(raw_items)} events into DB", flush=True)
-                    now = time.time()
-                    # if now - last_quota >= cfg["quota_refresh_seconds"]:
-                       # refresh_quota(force=True)
-                       # last_quota = now
-                    time.sleep(cfg["poll_interval_seconds"])
+                if mode == "subscribe":
+                    _collect_subscribe(client, cfg, channel)
+                else:
+                    _collect_rpop(client, cfg, channel)
             finally:
                 client.close()
                 print("collector: connection closed", flush=True)
@@ -435,6 +491,73 @@ def collect_forever():
             time.sleep(5)
 
 
+def _collect_subscribe(client, cfg, channel):
+    """Collector dùng pub/sub SUBSCRIBE — nhận event theo thời gian thực.
+    Dùng select() thay vì socket.settimeout() để tránh OSError trên makefile.
+    """
+    import select as _select
+
+    reply = client.subscribe(channel)
+    print(f"collector: SUBSCRIBE {channel!r} → {reply}", flush=True)
+
+    # Để socket ở blocking mode (không settimeout) — dùng select để poll
+    client.sock.settimeout(None)
+
+    batch = []
+    last_flush = time.time()
+    SELECT_TIMEOUT = 2.0  # giây
+
+    while True:
+        # Chờ có dữ liệu đến trong SELECT_TIMEOUT giây
+        readable, _, exceptional = _select.select(
+            [client.sock], [], [client.sock], SELECT_TIMEOUT
+        )
+
+        if exceptional:
+            raise EOFError("socket error detected via select")
+
+        if readable:
+            try:
+                data = client.read_message()
+            except EOFError:
+                raise  # kết nối đứt → vòng ngoài reconnect
+            except Exception as e:
+                print(f"collector: read_message error: {e}", file=sys.stderr, flush=True)
+                data = None
+
+            if data is not None:
+                # Lọc heartbeat {"support_refresh":true}
+                stripped = data.strip()
+                if stripped.startswith("{"):
+                    try:
+                        obj = json.loads(stripped)
+                        if isinstance(obj, dict) and set(obj.keys()) == {"support_refresh"}:
+                            continue  # bỏ qua heartbeat
+                    except Exception:
+                        pass
+                batch.append(data)
+
+        # Flush batch sau mỗi 50 item hoặc mỗi poll_interval giây
+        now = time.time()
+        if batch and (len(batch) >= 50 or now - last_flush >= cfg["poll_interval_seconds"]):
+            print(f"collector: flushing {len(batch)} event(s) from pub/sub", flush=True)
+            inserted = insert_usage(batch)
+            print(f"collector: inserted {inserted}/{len(batch)} events into DB", flush=True)
+            batch.clear()
+            last_flush = now
+
+
+def _collect_rpop(client, cfg, channel):
+    """Collector dùng RPOP (fallback cho queue mode)."""
+    while True:
+        raw_items = client.rpop(100, channel=channel)
+        if raw_items:
+            print(f"collector: got {len(raw_items)} item(s) via RPOP {channel!r}", flush=True)
+            inserted = insert_usage(raw_items)
+            print(f"collector: inserted {inserted}/{len(raw_items)} events into DB", flush=True)
+        time.sleep(cfg["poll_interval_seconds"])
+
+
 def range_bounds(name):
     now = dt.datetime.now(LOCAL_TZ)
     if name == "5h":