diff --git a/tools/pc-metrics/push-metrics-ble.py b/tools/pc-metrics/push-metrics-ble.py index aaff7e5..a968d32 100644 --- a/tools/pc-metrics/push-metrics-ble.py +++ b/tools/pc-metrics/push-metrics-ble.py @@ -1,11 +1,12 @@ import argparse import asyncio import ctypes -from datetime import date, timedelta +from datetime import date, datetime, timedelta import json import os from pathlib import Path import queue +import sqlite3 import subprocess import sys import threading @@ -24,8 +25,8 @@ DEFAULT_CHAR_UUID = "9f4af548-1a14-4fc1-8f5a-4b9f00000002" MAX_DISK_ROWS = 4 MAX_DISK_TEXT = 124 CODEX_DAILY_BAR_COUNT = 14 -CODEX_REFRESH_SECONDS = 60.0 -CODEX_USAGE_REFRESH_SECONDS = 300.0 +CODEX_REFRESH_SECONDS = 300.0 +CCSWITCH_CODEX_REFRESH_SECONDS = 60.0 ZHIPU_REFRESH_SECONDS = 60.0 ZHIPU_USAGE_REFRESH_SECONDS = 600.0 ZHIPU_DAILY_BAR_COUNT = 14 @@ -45,6 +46,25 @@ def clamp_kbps(value: float) -> int: return max(0, min(999999, int(round(value)))) +def tokens_to_m10(tokens) -> int: + if not isinstance(tokens, (int, float)) or tokens < 0: + return -1 + return max(0, min(2147483647, int(round(tokens / 100000)))) + + +def daily_bars_from_token_map(tokens_by_day: dict[str, int], count: int) -> tuple[int, str]: + today = date.today() + wanted_days = [(today - timedelta(days=offset)).isoformat() for offset in range(count - 1, -1, -1)] + values = [tokens_by_day.get(day, 0) for day in wanted_days] + max_value = max(values) if values else 0 + if max_value > 0: + percents = [max(0, min(100, int(round(value * 100 / max_value)))) for value in values] + else: + percents = [0 for _ in values] + today_tokens_m10 = tokens_to_m10(tokens_by_day.get(today.isoformat(), 0)) + return today_tokens_m10, ":".join(str(value) for value in percents) + + class _PdhCounterValue(ctypes.Structure): _fields_ = [ ("status", wintypes.LONG), @@ -421,9 +441,6 @@ class CodexRateLimitReader: self.refresh_seconds = refresh_seconds self.next_refresh = 0.0 self.cached = (-1, -1, -1, -1, -1) - self.usage_refresh_seconds = CODEX_USAGE_REFRESH_SECONDS - self.next_usage_refresh = 0.0 - self.usage_cached = (-1, "") self.process: subprocess.Popen | None = None self.responses: queue.Queue[dict] = queue.Queue() self.next_id = 1 @@ -579,60 +596,144 @@ class CodexRateLimitReader: ) return self.cached - @staticmethod - def _tokens_to_m10(tokens) -> int: - if not isinstance(tokens, (int, float)) or tokens < 0: - return -1 - return max(0, min(2147483647, int(round(tokens / 100000)))) +class CcswitchCodexUsageReader: + DEDUP_WINDOW_SECONDS = 10 * 60 + + def __init__(self, refresh_seconds: float = CCSWITCH_CODEX_REFRESH_SECONDS) -> None: + self.refresh_seconds = refresh_seconds + self.next_refresh = 0.0 + self.cached = (-1, "") @staticmethod - def _daily_bars_from_buckets(buckets: list[dict]) -> tuple[int, str]: - today = date.today() - wanted_days = [(today - timedelta(days=offset)).isoformat() for offset in range(CODEX_DAILY_BAR_COUNT - 1, -1, -1)] + def _candidate_paths() -> list[Path]: + paths: list[Path] = [] + override = os.environ.get("CCSWITCH_DB_PATH", "").strip() + if override: + paths.append(Path(override).expanduser()) + + home = Path.home() + paths.append(home / ".cc-switch" / "cc-switch.db") + + appdata = os.environ.get("APPDATA") + if appdata: + paths.append(Path(appdata) / "com.ccswitch.desktop" / "cc-switch.db") + + local_appdata = os.environ.get("LOCALAPPDATA") + if local_appdata: + paths.append(Path(local_appdata) / "com.ccswitch.desktop" / "cc-switch.db") + + unique: list[Path] = [] + seen: set[str] = set() + for path in paths: + key = str(path).lower() + if key not in seen: + seen.add(key) + unique.append(path) + return unique + + @classmethod + def _find_db(cls) -> Path | None: + for path in cls._candidate_paths(): + if path.is_file(): + return path + return None + + @staticmethod + def _connect_readonly(path: Path) -> sqlite3.Connection: + sqlite_path = urllib.parse.quote(str(path).replace("\\", "/"), safe="/:") + con = sqlite3.connect(f"file:{sqlite_path}?mode=ro", uri=True, timeout=1.0) + con.row_factory = sqlite3.Row + return con + + @classmethod + def _effective_filter_sql(cls) -> str: + return f""" + NOT ( + COALESCE(l.data_source, 'proxy') IN ('session_log', 'codex_session', 'gemini_session', 'opencode_session') + AND EXISTS ( + SELECT 1 + FROM proxy_request_logs proxy_dedup + WHERE COALESCE(proxy_dedup.data_source, 'proxy') = 'proxy' + AND proxy_dedup.app_type = l.app_type + AND proxy_dedup.status_code >= 200 + AND proxy_dedup.status_code < 300 + AND proxy_dedup.input_tokens = l.input_tokens + AND proxy_dedup.output_tokens = l.output_tokens + AND proxy_dedup.cache_read_tokens = l.cache_read_tokens + AND ( + proxy_dedup.cache_creation_tokens = l.cache_creation_tokens + OR ( + l.cache_creation_tokens = 0 + AND COALESCE(l.data_source, 'proxy') IN ('codex_session', 'gemini_session', 'opencode_session') + ) + ) + AND proxy_dedup.created_at BETWEEN + l.created_at - {cls.DEDUP_WINDOW_SECONDS} + AND l.created_at + {cls.DEDUP_WINDOW_SECONDS} + AND ( + LOWER(proxy_dedup.model) = LOWER(l.model) + OR LOWER(proxy_dedup.model) = 'unknown' + OR LOWER(l.model) = 'unknown' + ) + ) + ) + """ + + def _query_daily_tokens(self, db_path: Path) -> dict[str, int]: + start_day = date.today() - timedelta(days=CODEX_DAILY_BAR_COUNT - 1) + start_ts = int(datetime.combine(start_day, datetime.min.time()).timestamp()) + real_total_sql = """ + ( + CASE + WHEN COALESCE(l.input_tokens, 0) >= COALESCE(l.cache_read_tokens, 0) + THEN COALESCE(l.input_tokens, 0) - COALESCE(l.cache_read_tokens, 0) + ELSE COALESCE(l.input_tokens, 0) + END + + COALESCE(l.output_tokens, 0) + + COALESCE(l.cache_creation_tokens, 0) + + COALESCE(l.cache_read_tokens, 0) + ) + """ + sql = f""" + SELECT date(l.created_at, 'unixepoch', 'localtime') AS usage_day, + SUM({real_total_sql}) AS real_total_tokens + FROM proxy_request_logs l + WHERE l.app_type = 'codex' + AND l.status_code >= 200 + AND l.status_code < 300 + AND l.created_at >= ? + AND {self._effective_filter_sql()} + GROUP BY usage_day + """ + tokens_by_day: dict[str, int] = {} - - for bucket in buckets: - if not isinstance(bucket, dict): - continue - start_date = bucket.get("startDate") - tokens = bucket.get("tokens") - if isinstance(start_date, str) and isinstance(tokens, (int, float)) and tokens >= 0: - tokens_by_day[start_date] = int(tokens) - - values = [tokens_by_day.get(day, 0) for day in wanted_days] - max_value = max(values) if values else 0 - if max_value > 0: - percents = [max(0, min(100, int(round(value * 100 / max_value)))) for value in values] - else: - percents = [0 for _ in values] - - today_tokens_m10 = CodexRateLimitReader._tokens_to_m10(tokens_by_day.get(today.isoformat(), 0)) - return today_tokens_m10, ":".join(str(value) for value in percents) + with self._connect_readonly(db_path) as con: + for row in con.execute(sql, (start_ts,)): + usage_day = row["usage_day"] + tokens = row["real_total_tokens"] + if isinstance(usage_day, str) and isinstance(tokens, (int, float)) and tokens >= 0: + tokens_by_day[usage_day] = int(tokens) + return tokens_by_day def read_usage_snapshot(self) -> tuple[int, str]: now = time.monotonic() - if now < self.next_usage_refresh: - return self.usage_cached + if now < self.next_refresh: + return self.cached - if not self._ensure_started(): - self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds) - return self.usage_cached + db_path = self._find_db() + if db_path is None: + self.next_refresh = now + min(15.0, self.refresh_seconds) + return self.cached - request_id = self._send("account/usage/read") - if request_id is None: - self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds) - return self.usage_cached + try: + tokens_by_day = self._query_daily_tokens(db_path) + except sqlite3.Error: + self.next_refresh = now + min(15.0, self.refresh_seconds) + return self.cached - message = self._wait_for(request_id, 8.0) - result = message.get("result") if message else None - buckets = result.get("dailyUsageBuckets") if isinstance(result, dict) else None - if not isinstance(buckets, list) or not buckets: - self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds) - return self.usage_cached - - self.next_usage_refresh = now + self.usage_refresh_seconds - self.usage_cached = self._daily_bars_from_buckets(buckets) - return self.usage_cached + self.cached = daily_bars_from_token_map(tokens_by_day, CODEX_DAILY_BAR_COUNT) + self.next_refresh = now + self.refresh_seconds + return self.cached class ZhipuStatusReader: @@ -824,6 +925,7 @@ class MetricsSampler: self.cpu_frequency = CpuFrequencyReader() self.disk_rates = LogicalDiskRateReader(self._disk_labels()) self.codex = CodexRateLimitReader() + self.codex_usage = CcswitchCodexUsageReader() self.zhipu = ZhipuStatusReader() psutil.cpu_percent(interval=None) @@ -927,7 +1029,7 @@ class MetricsSampler: disk_count, disk_percent, disk_text = self._disk_summary(disk_rates) codex_5h, codex_week, codex_5h_reset, codex_week_reset, codex_credits = self.codex.read() - codex_today_tokens_m10, codex_daily_bars = self.codex.read_usage_snapshot() + codex_today_tokens_m10, codex_daily_bars = self.codex_usage.read_usage_snapshot() zhipu_status, zhipu_5h, zhipu_week, zhipu_5h_reset, zhipu_week_reset, zhipu_today_tokens_m10, zhipu_daily_bars = self.zhipu.read() return (