843 lines
30 KiB
Python
843 lines
30 KiB
Python
import argparse
|
|
import asyncio
|
|
import ctypes
|
|
from datetime import date, timedelta
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import queue
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from ctypes import wintypes
|
|
|
|
import psutil
|
|
from bleak import BleakClient, BleakScanner
|
|
|
|
|
|
DEFAULT_DEVICE_NAME = "RLCD-Metrics"
|
|
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
|
|
ZHIPU_REFRESH_SECONDS = 60.0
|
|
|
|
|
|
def clamp_percent(value: float) -> int:
|
|
if value != value:
|
|
return 0
|
|
return max(0, min(100, int(round(value))))
|
|
|
|
|
|
def mb(value: float) -> int:
|
|
return int(round(value / 1024 / 1024))
|
|
|
|
|
|
def clamp_kbps(value: float) -> int:
|
|
return max(0, min(999999, int(round(value))))
|
|
|
|
|
|
class _PdhCounterValue(ctypes.Structure):
|
|
_fields_ = [
|
|
("status", wintypes.LONG),
|
|
("double_value", ctypes.c_double),
|
|
]
|
|
|
|
|
|
class CpuPercentReader:
|
|
def __init__(self) -> None:
|
|
self.pdh = None
|
|
self.query = ctypes.c_void_p()
|
|
self.counter = ctypes.c_void_p()
|
|
self.available = False
|
|
|
|
if sys.platform != "win32":
|
|
return
|
|
|
|
try:
|
|
self.pdh = ctypes.WinDLL("pdh")
|
|
self.pdh.PdhOpenQueryW.restype = wintypes.LONG
|
|
self.pdh.PdhOpenQueryW.argtypes = [
|
|
wintypes.LPCWSTR,
|
|
ctypes.c_void_p,
|
|
ctypes.POINTER(ctypes.c_void_p),
|
|
]
|
|
self.pdh.PdhCollectQueryData.restype = wintypes.LONG
|
|
self.pdh.PdhCollectQueryData.argtypes = [ctypes.c_void_p]
|
|
self.pdh.PdhGetFormattedCounterValue.restype = wintypes.LONG
|
|
self.pdh.PdhGetFormattedCounterValue.argtypes = [
|
|
ctypes.c_void_p,
|
|
wintypes.DWORD,
|
|
ctypes.POINTER(wintypes.DWORD),
|
|
ctypes.POINTER(_PdhCounterValue),
|
|
]
|
|
|
|
add_counter = getattr(self.pdh, "PdhAddEnglishCounterW", self.pdh.PdhAddCounterW)
|
|
add_counter.restype = wintypes.LONG
|
|
add_counter.argtypes = [
|
|
ctypes.c_void_p,
|
|
wintypes.LPCWSTR,
|
|
ctypes.c_void_p,
|
|
ctypes.POINTER(ctypes.c_void_p),
|
|
]
|
|
|
|
if self.pdh.PdhOpenQueryW(None, None, ctypes.byref(self.query)) != 0:
|
|
return
|
|
if add_counter(
|
|
self.query,
|
|
r"\Processor Information(_Total)\% Processor Utility",
|
|
None,
|
|
ctypes.byref(self.counter),
|
|
) != 0:
|
|
return
|
|
|
|
self.pdh.PdhCollectQueryData(self.query)
|
|
self.available = True
|
|
except OSError:
|
|
self.available = False
|
|
|
|
def read_percent(self) -> int:
|
|
if self.available and self.pdh is not None:
|
|
value = _PdhCounterValue()
|
|
counter_type = wintypes.DWORD()
|
|
if self.pdh.PdhCollectQueryData(self.query) == 0:
|
|
rc = self.pdh.PdhGetFormattedCounterValue(
|
|
self.counter,
|
|
0x00000200,
|
|
ctypes.byref(counter_type),
|
|
ctypes.byref(value),
|
|
)
|
|
if rc == 0 and value.status == 0 and value.double_value >= 0:
|
|
return clamp_percent(value.double_value)
|
|
|
|
return clamp_percent(psutil.cpu_percent(interval=None))
|
|
|
|
|
|
class CpuFrequencyReader:
|
|
def __init__(self) -> None:
|
|
self.pdh = None
|
|
self.query = ctypes.c_void_p()
|
|
self.total_counter = ctypes.c_void_p()
|
|
self.core_counters: list[tuple[int, ctypes.c_void_p]] = []
|
|
self.base_mhz = self._base_mhz()
|
|
self.available = False
|
|
|
|
if sys.platform != "win32":
|
|
return
|
|
|
|
try:
|
|
self.pdh = ctypes.WinDLL("pdh")
|
|
self.pdh.PdhOpenQueryW.restype = wintypes.LONG
|
|
self.pdh.PdhOpenQueryW.argtypes = [
|
|
wintypes.LPCWSTR,
|
|
ctypes.c_void_p,
|
|
ctypes.POINTER(ctypes.c_void_p),
|
|
]
|
|
self.pdh.PdhCollectQueryData.restype = wintypes.LONG
|
|
self.pdh.PdhCollectQueryData.argtypes = [ctypes.c_void_p]
|
|
self.pdh.PdhGetFormattedCounterValue.restype = wintypes.LONG
|
|
self.pdh.PdhGetFormattedCounterValue.argtypes = [
|
|
ctypes.c_void_p,
|
|
wintypes.DWORD,
|
|
ctypes.POINTER(wintypes.DWORD),
|
|
ctypes.POINTER(_PdhCounterValue),
|
|
]
|
|
|
|
add_counter = getattr(self.pdh, "PdhAddEnglishCounterW", self.pdh.PdhAddCounterW)
|
|
add_counter.restype = wintypes.LONG
|
|
add_counter.argtypes = [
|
|
ctypes.c_void_p,
|
|
wintypes.LPCWSTR,
|
|
ctypes.c_void_p,
|
|
ctypes.POINTER(ctypes.c_void_p),
|
|
]
|
|
|
|
if self.pdh.PdhOpenQueryW(None, None, ctypes.byref(self.query)) != 0:
|
|
return
|
|
if add_counter(
|
|
self.query,
|
|
r"\Processor Information(_Total)\% Processor Performance",
|
|
None,
|
|
ctypes.byref(self.total_counter),
|
|
) != 0:
|
|
self.total_counter = ctypes.c_void_p()
|
|
|
|
logical_count = psutil.cpu_count(logical=True) or 0
|
|
for index in range(logical_count):
|
|
counter = ctypes.c_void_p()
|
|
path = rf"\Processor Information(0,{index})\% Processor Performance"
|
|
if add_counter(self.query, path, None, ctypes.byref(counter)) == 0:
|
|
self.core_counters.append((index, counter))
|
|
|
|
if not self.total_counter and not self.core_counters:
|
|
return
|
|
|
|
self.pdh.PdhCollectQueryData(self.query)
|
|
self.available = True
|
|
except OSError:
|
|
self.available = False
|
|
|
|
@staticmethod
|
|
def _base_mhz() -> int:
|
|
freq = psutil.cpu_freq()
|
|
if freq and freq.max:
|
|
return int(round(freq.max))
|
|
if freq and freq.current:
|
|
return int(round(freq.current))
|
|
return 0
|
|
|
|
@staticmethod
|
|
def _format_core_mhz(index: int, mhz: int) -> str:
|
|
ghz10 = max(0, int(round(mhz / 100)))
|
|
return f"{index:02d}:{ghz10 // 10}.{ghz10 % 10}"
|
|
|
|
def _read_counter(self, counter: ctypes.c_void_p) -> float:
|
|
if self.pdh is None or not counter:
|
|
return -1.0
|
|
|
|
value = _PdhCounterValue()
|
|
counter_type = wintypes.DWORD()
|
|
rc = self.pdh.PdhGetFormattedCounterValue(
|
|
counter,
|
|
0x00000200,
|
|
ctypes.byref(counter_type),
|
|
ctypes.byref(value),
|
|
)
|
|
if rc != 0 or value.status != 0 or value.double_value < 0:
|
|
return -1.0
|
|
return value.double_value
|
|
|
|
def _perf_to_mhz(self, performance_percent: float) -> int:
|
|
if self.base_mhz <= 0 or performance_percent < 0:
|
|
return 0
|
|
return int(round(self.base_mhz * performance_percent / 100))
|
|
|
|
def read(self) -> tuple[int, str]:
|
|
if self.available and self.pdh is not None and self.base_mhz > 0:
|
|
if self.pdh.PdhCollectQueryData(self.query) == 0:
|
|
total_mhz = self._perf_to_mhz(self._read_counter(self.total_counter))
|
|
core_values: list[tuple[int, int]] = []
|
|
for index, counter in self.core_counters:
|
|
mhz = self._perf_to_mhz(self._read_counter(counter))
|
|
if mhz > 0:
|
|
core_values.append((index, mhz))
|
|
|
|
core_values.sort(key=lambda item: (-item[1], item[0]))
|
|
top_text = " ".join(
|
|
self._format_core_mhz(index, mhz)
|
|
for index, mhz in core_values[:4]
|
|
)
|
|
if total_mhz <= 0 and core_values:
|
|
total_mhz = core_values[0][1]
|
|
if total_mhz > 0:
|
|
return total_mhz, top_text
|
|
|
|
freq = psutil.cpu_freq()
|
|
return (int(round(freq.current)) if freq and freq.current else 0), ""
|
|
|
|
def read_mhz(self) -> int:
|
|
return self.read()[0]
|
|
|
|
|
|
class LogicalDiskRateReader:
|
|
def __init__(self, labels: list[str]) -> None:
|
|
self.pdh = None
|
|
self.query = ctypes.c_void_p()
|
|
self.counters: dict[str, tuple[ctypes.c_void_p, ctypes.c_void_p]] = {}
|
|
self.available = False
|
|
|
|
if sys.platform != "win32" or not labels:
|
|
return
|
|
|
|
try:
|
|
self.pdh = ctypes.WinDLL("pdh")
|
|
self.pdh.PdhOpenQueryW.restype = wintypes.LONG
|
|
self.pdh.PdhOpenQueryW.argtypes = [
|
|
wintypes.LPCWSTR,
|
|
ctypes.c_void_p,
|
|
ctypes.POINTER(ctypes.c_void_p),
|
|
]
|
|
self.pdh.PdhCollectQueryData.restype = wintypes.LONG
|
|
self.pdh.PdhCollectQueryData.argtypes = [ctypes.c_void_p]
|
|
self.pdh.PdhGetFormattedCounterValue.restype = wintypes.LONG
|
|
self.pdh.PdhGetFormattedCounterValue.argtypes = [
|
|
ctypes.c_void_p,
|
|
wintypes.DWORD,
|
|
ctypes.POINTER(wintypes.DWORD),
|
|
ctypes.POINTER(_PdhCounterValue),
|
|
]
|
|
|
|
add_counter = getattr(self.pdh, "PdhAddEnglishCounterW", self.pdh.PdhAddCounterW)
|
|
add_counter.restype = wintypes.LONG
|
|
add_counter.argtypes = [
|
|
ctypes.c_void_p,
|
|
wintypes.LPCWSTR,
|
|
ctypes.c_void_p,
|
|
ctypes.POINTER(ctypes.c_void_p),
|
|
]
|
|
|
|
if self.pdh.PdhOpenQueryW(None, None, ctypes.byref(self.query)) != 0:
|
|
return
|
|
|
|
for label in labels:
|
|
read_counter = ctypes.c_void_p()
|
|
write_counter = ctypes.c_void_p()
|
|
read_path = rf"\LogicalDisk({label}:)\Disk Read Bytes/sec"
|
|
write_path = rf"\LogicalDisk({label}:)\Disk Write Bytes/sec"
|
|
if add_counter(self.query, read_path, None, ctypes.byref(read_counter)) != 0:
|
|
continue
|
|
if add_counter(self.query, write_path, None, ctypes.byref(write_counter)) != 0:
|
|
continue
|
|
self.counters[label] = (read_counter, write_counter)
|
|
|
|
if self.counters:
|
|
self.pdh.PdhCollectQueryData(self.query)
|
|
self.available = True
|
|
except OSError:
|
|
self.available = False
|
|
|
|
def _read_kbps(self, counter: ctypes.c_void_p) -> int:
|
|
if self.pdh is None:
|
|
return -1
|
|
|
|
value = _PdhCounterValue()
|
|
counter_type = wintypes.DWORD()
|
|
rc = self.pdh.PdhGetFormattedCounterValue(
|
|
counter,
|
|
0x00000200,
|
|
ctypes.byref(counter_type),
|
|
ctypes.byref(value),
|
|
)
|
|
if rc != 0 or value.status != 0 or value.double_value < 0:
|
|
return -1
|
|
return clamp_kbps(value.double_value / 1024)
|
|
|
|
def read_rates(self) -> dict[str, tuple[int, int]]:
|
|
if not self.available or self.pdh is None:
|
|
return {}
|
|
|
|
if self.pdh.PdhCollectQueryData(self.query) != 0:
|
|
return {}
|
|
|
|
rates: dict[str, tuple[int, int]] = {}
|
|
for label, (read_counter, write_counter) in self.counters.items():
|
|
rates[label] = (self._read_kbps(read_counter), self._read_kbps(write_counter))
|
|
return rates
|
|
|
|
|
|
class CodexRateLimitReader:
|
|
def __init__(self, refresh_seconds: float = CODEX_REFRESH_SECONDS) -> None:
|
|
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
|
|
self.initialized = False
|
|
|
|
def _find_codex_exe(self) -> str | None:
|
|
local_app_data = os.environ.get("LOCALAPPDATA")
|
|
if local_app_data:
|
|
root = Path(local_app_data) / "OpenAI" / "Codex" / "bin"
|
|
candidates = sorted(root.glob("*/codex.exe"), key=lambda p: p.stat().st_mtime, reverse=True)
|
|
if candidates:
|
|
return str(candidates[0])
|
|
return "codex.exe"
|
|
|
|
def _reader_loop(self, stdout) -> None:
|
|
for line in stdout:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
self.responses.put(json.loads(line))
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
def _stop(self) -> None:
|
|
if self.process is not None:
|
|
try:
|
|
self.process.kill()
|
|
except OSError:
|
|
pass
|
|
self.process = None
|
|
self.initialized = False
|
|
|
|
def _send(self, method: str, params: dict | None = None, expect_response: bool = True) -> int | None:
|
|
if self.process is None or self.process.stdin is None:
|
|
return None
|
|
|
|
request_id = None
|
|
message: dict = {"method": method}
|
|
if params is not None:
|
|
message["params"] = params
|
|
if expect_response:
|
|
request_id = self.next_id
|
|
self.next_id += 1
|
|
message["id"] = request_id
|
|
|
|
self.process.stdin.write(json.dumps(message, separators=(",", ":")) + "\n")
|
|
self.process.stdin.flush()
|
|
return request_id
|
|
|
|
def _wait_for(self, request_id: int, timeout: float) -> dict | None:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
try:
|
|
message = self.responses.get(timeout=max(0.05, deadline - time.monotonic()))
|
|
except queue.Empty:
|
|
return None
|
|
if message.get("id") == request_id:
|
|
return message
|
|
return None
|
|
|
|
def _ensure_started(self) -> bool:
|
|
if self.process is not None and self.process.poll() is None and self.initialized:
|
|
return True
|
|
|
|
self._stop()
|
|
codex_exe = self._find_codex_exe()
|
|
if not codex_exe:
|
|
return False
|
|
|
|
creationflags = subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
|
|
try:
|
|
self.process = subprocess.Popen(
|
|
[codex_exe, "app-server", "--stdio"],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.DEVNULL,
|
|
text=True,
|
|
encoding="utf-8",
|
|
bufsize=1,
|
|
creationflags=creationflags,
|
|
)
|
|
except OSError:
|
|
self.process = None
|
|
return False
|
|
|
|
assert self.process.stdout is not None
|
|
threading.Thread(target=self._reader_loop, args=(self.process.stdout,), daemon=True).start()
|
|
|
|
init_id = self._send(
|
|
"initialize",
|
|
{
|
|
"clientInfo": {
|
|
"name": "rlcd_metrics_bridge",
|
|
"title": "RLCD Metrics Bridge",
|
|
"version": "0.1.0",
|
|
}
|
|
},
|
|
)
|
|
if init_id is None or self._wait_for(init_id, 5.0) is None:
|
|
self._stop()
|
|
return False
|
|
|
|
self._send("initialized", {}, expect_response=False)
|
|
self.initialized = True
|
|
return True
|
|
|
|
@staticmethod
|
|
def _mins_until(epoch_seconds) -> int:
|
|
if not isinstance(epoch_seconds, (int, float)):
|
|
return -1
|
|
return max(0, min(99999, int(round((epoch_seconds - time.time()) / 60))))
|
|
|
|
@staticmethod
|
|
def _read_percent(value) -> int:
|
|
if not isinstance(value, (int, float)):
|
|
return -1
|
|
return clamp_percent(value)
|
|
|
|
def read(self) -> tuple[int, int, int, int, int]:
|
|
now = time.monotonic()
|
|
if now < self.next_refresh:
|
|
return self.cached
|
|
|
|
self.next_refresh = now + self.refresh_seconds
|
|
if not self._ensure_started():
|
|
return self.cached
|
|
|
|
request_id = self._send("account/rateLimits/read")
|
|
if request_id is None:
|
|
return self.cached
|
|
|
|
message = self._wait_for(request_id, 8.0)
|
|
result = message.get("result") if message else None
|
|
rate_limits = result.get("rateLimits") if isinstance(result, dict) else None
|
|
if not isinstance(rate_limits, dict):
|
|
self._stop()
|
|
return self.cached
|
|
|
|
primary = rate_limits.get("primary") or {}
|
|
secondary = rate_limits.get("secondary") or {}
|
|
reset_credits = result.get("rateLimitResetCredits") or {}
|
|
|
|
self.cached = (
|
|
self._read_percent(primary.get("usedPercent")),
|
|
self._read_percent(secondary.get("usedPercent")),
|
|
self._mins_until(primary.get("resetsAt")),
|
|
self._mins_until(secondary.get("resetsAt")),
|
|
int(reset_credits.get("availableCount", -1)) if isinstance(reset_credits, dict) else -1,
|
|
)
|
|
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))))
|
|
|
|
@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)]
|
|
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(), -1))
|
|
return today_tokens_m10, ":".join(str(value) for value in percents)
|
|
|
|
def read_usage_snapshot(self) -> tuple[int, str]:
|
|
now = time.monotonic()
|
|
if now < self.next_usage_refresh:
|
|
return self.usage_cached
|
|
|
|
self.next_usage_refresh = now + self.usage_refresh_seconds
|
|
if not self._ensure_started():
|
|
return self.usage_cached
|
|
|
|
request_id = self._send("account/usage/read")
|
|
if request_id is None:
|
|
return self.usage_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:
|
|
return self.usage_cached
|
|
|
|
self.usage_cached = self._daily_bars_from_buckets(buckets)
|
|
return self.usage_cached
|
|
|
|
|
|
class ZhipuStatusReader:
|
|
def __init__(self, refresh_seconds: float = ZHIPU_REFRESH_SECONDS) -> None:
|
|
self.refresh_seconds = refresh_seconds
|
|
self.next_refresh = 0.0
|
|
self.cached = (0, -1, -1, -1, -1)
|
|
|
|
@staticmethod
|
|
def _quota_url() -> str:
|
|
base_url = os.environ.get("ZHIPU_BASE_URL", "https://open.bigmodel.cn").strip()
|
|
if "bigmodel.cn" in base_url.lower():
|
|
host = "https://open.bigmodel.cn"
|
|
else:
|
|
host = "https://api.z.ai"
|
|
return f"{host}/api/monitor/usage/quota/limit"
|
|
|
|
@staticmethod
|
|
def _mins_until(epoch_ms) -> int:
|
|
if not isinstance(epoch_ms, (int, float)) or epoch_ms <= 0:
|
|
return -1
|
|
return max(0, min(99999, int(round((epoch_ms / 1000 - time.time()) / 60))))
|
|
|
|
@staticmethod
|
|
def _read_percent(value) -> int:
|
|
if isinstance(value, (int, float)):
|
|
return clamp_percent(value)
|
|
try:
|
|
return clamp_percent(float(value))
|
|
except (TypeError, ValueError):
|
|
return -1
|
|
|
|
def _query(self, api_key: str) -> tuple[int, int, int, int, int]:
|
|
request = urllib.request.Request(
|
|
self._quota_url(),
|
|
headers={
|
|
"Authorization": api_key,
|
|
"Content-Type": "application/json",
|
|
"Accept-Language": "en-US,en",
|
|
},
|
|
method="GET",
|
|
)
|
|
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=15) as response:
|
|
body = response.read().decode("utf-8", errors="replace")
|
|
except urllib.error.HTTPError as exc:
|
|
if exc.code in (401, 403):
|
|
return (2, -1, -1, -1, -1)
|
|
return (2, -1, -1, -1, -1)
|
|
except OSError:
|
|
return (2, -1, -1, -1, -1)
|
|
|
|
try:
|
|
payload = json.loads(body)
|
|
except json.JSONDecodeError:
|
|
return (2, -1, -1, -1, -1)
|
|
|
|
if payload.get("success") is False:
|
|
return (2, -1, -1, -1, -1)
|
|
|
|
data = payload.get("data")
|
|
limits = data.get("limits") if isinstance(data, dict) else None
|
|
if not isinstance(limits, list):
|
|
return (2, -1, -1, -1, -1)
|
|
|
|
five_percent = -1
|
|
week_percent = -1
|
|
five_reset = -1
|
|
week_reset = -1
|
|
|
|
for item in limits:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
if str(item.get("type", "")).upper() != "TOKENS_LIMIT":
|
|
continue
|
|
unit = item.get("unit")
|
|
percent = self._read_percent(item.get("percentage"))
|
|
if percent < 0:
|
|
continue
|
|
reset_mins = self._mins_until(item.get("nextResetTime"))
|
|
if unit == 3 and five_percent < 0:
|
|
five_percent = percent
|
|
five_reset = reset_mins
|
|
elif unit == 6 and week_percent < 0:
|
|
week_percent = percent
|
|
week_reset = reset_mins
|
|
|
|
return (1, five_percent, week_percent, five_reset, week_reset)
|
|
|
|
def read(self) -> tuple[int, int, int, int, int]:
|
|
api_key = os.environ.get("ZHIPU_API_KEY", "").strip()
|
|
if not api_key:
|
|
self.cached = (0, -1, -1, -1, -1)
|
|
self.next_refresh = time.monotonic() + self.refresh_seconds
|
|
return self.cached
|
|
|
|
now = time.monotonic()
|
|
if now < self.next_refresh:
|
|
return self.cached
|
|
|
|
self.next_refresh = now + self.refresh_seconds
|
|
self.cached = self._query(api_key)
|
|
return self.cached
|
|
|
|
|
|
class MetricsSampler:
|
|
def __init__(self, disk_path: str) -> None:
|
|
self.disk_path = disk_path
|
|
self.last_net = psutil.net_io_counters()
|
|
self.last_disk = psutil.disk_io_counters()
|
|
self.last_ts = time.monotonic()
|
|
self.cpu_percent = CpuPercentReader()
|
|
self.cpu_frequency = CpuFrequencyReader()
|
|
self.disk_rates = LogicalDiskRateReader(self._disk_labels())
|
|
self.codex = CodexRateLimitReader()
|
|
self.zhipu = ZhipuStatusReader()
|
|
psutil.cpu_percent(interval=None)
|
|
|
|
def _disk_labels(self) -> list[str]:
|
|
labels = []
|
|
seen = set()
|
|
|
|
for part in psutil.disk_partitions(all=False):
|
|
device = (part.device or part.mountpoint or "").upper()
|
|
if not device or device in seen:
|
|
continue
|
|
seen.add(device)
|
|
|
|
opts = (part.opts or "").lower()
|
|
if "cdrom" in opts or not part.fstype:
|
|
continue
|
|
|
|
label = device.replace("\\", "").replace("/", "")
|
|
if label.endswith(":"):
|
|
label = label[:-1]
|
|
if label:
|
|
labels.append(label)
|
|
|
|
return labels
|
|
|
|
def _disk_summary(self, rates_by_label: dict[str, tuple[int, int]]) -> tuple[int, int, str]:
|
|
total = 0
|
|
used = 0
|
|
parts = []
|
|
seen = set()
|
|
|
|
for part in psutil.disk_partitions(all=False):
|
|
device = (part.device or part.mountpoint or "").upper()
|
|
if not device or device in seen:
|
|
continue
|
|
seen.add(device)
|
|
|
|
opts = (part.opts or "").lower()
|
|
if "cdrom" in opts or not part.fstype:
|
|
continue
|
|
|
|
try:
|
|
usage = psutil.disk_usage(part.mountpoint)
|
|
except OSError:
|
|
continue
|
|
|
|
total += usage.total
|
|
used += usage.used
|
|
label = device.replace("\\", "").replace("/", "")
|
|
if label.endswith(":"):
|
|
label = label[:-1]
|
|
read_kbps, write_kbps = rates_by_label.get(label, (-1, -1))
|
|
parts.append(f"{label}:{clamp_percent(usage.percent)}:{read_kbps}:{write_kbps}")
|
|
|
|
if total <= 0:
|
|
try:
|
|
usage = psutil.disk_usage(self.disk_path)
|
|
return 1, clamp_percent(usage.percent), f"C:{clamp_percent(usage.percent)}:-1:-1"
|
|
except OSError:
|
|
return 0, 0, ""
|
|
|
|
disk_percent = clamp_percent((used / total) * 100)
|
|
selected = []
|
|
text_len = 0
|
|
for part in parts[:MAX_DISK_ROWS]:
|
|
extra = len(part) + (1 if selected else 0)
|
|
if text_len + extra > MAX_DISK_TEXT:
|
|
break
|
|
selected.append(part)
|
|
text_len += extra
|
|
text = " ".join(selected)
|
|
return len(parts), disk_percent, text
|
|
|
|
def sample_line(self) -> str:
|
|
now = time.monotonic()
|
|
elapsed = max(0.001, now - self.last_ts)
|
|
|
|
net = psutil.net_io_counters()
|
|
rx_kbps = int(max(0, round((net.bytes_recv - self.last_net.bytes_recv) / elapsed / 1024)))
|
|
tx_kbps = int(max(0, round((net.bytes_sent - self.last_net.bytes_sent) / elapsed / 1024)))
|
|
self.last_net = net
|
|
|
|
disk_io = psutil.disk_io_counters()
|
|
if disk_io is None or self.last_disk is None:
|
|
disk_read_kbps = 0
|
|
disk_write_kbps = 0
|
|
else:
|
|
disk_read_kbps = int(max(0, round((disk_io.read_bytes - self.last_disk.read_bytes) / elapsed / 1024)))
|
|
disk_write_kbps = int(max(0, round((disk_io.write_bytes - self.last_disk.write_bytes) / elapsed / 1024)))
|
|
self.last_disk = disk_io
|
|
self.last_ts = now
|
|
disk_rates = self.disk_rates.read_rates()
|
|
|
|
cpu = self.cpu_percent.read_percent()
|
|
cpu_mhz, cpu_core_text = self.cpu_frequency.read()
|
|
|
|
mem_info = psutil.virtual_memory()
|
|
mem_percent = clamp_percent(mem_info.percent)
|
|
mem_used_mb = mb(mem_info.used)
|
|
mem_total_mb = mb(mem_info.total)
|
|
|
|
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()
|
|
zhipu_status, zhipu_5h, zhipu_week, zhipu_5h_reset, zhipu_week_reset = self.zhipu.read()
|
|
|
|
return (
|
|
f"SYS7,{cpu},{cpu_mhz},{mem_percent},{mem_used_mb},{mem_total_mb},"
|
|
f"{disk_count},{disk_percent},{disk_read_kbps},{disk_write_kbps},"
|
|
f"{rx_kbps},{tx_kbps},"
|
|
f"{codex_5h},{codex_week},{codex_5h_reset},{codex_week_reset},{codex_credits},{codex_today_tokens_m10},"
|
|
f"{codex_daily_bars},"
|
|
f"{zhipu_status},{zhipu_5h},{zhipu_week},{zhipu_5h_reset},{zhipu_week_reset},{cpu_core_text},{disk_text}"
|
|
)
|
|
|
|
|
|
async def find_device(name: str, address: str | None, timeout: float):
|
|
if address:
|
|
return await BleakScanner.find_device_by_address(address, timeout=timeout)
|
|
|
|
def match(device, adv):
|
|
return device.name == name or adv.local_name == name
|
|
|
|
return await BleakScanner.find_device_by_filter(match, timeout=timeout)
|
|
|
|
|
|
async def push_loop(args) -> int:
|
|
sampler = MetricsSampler(args.disk)
|
|
|
|
while True:
|
|
print(f"Scanning for {args.name}...", flush=True)
|
|
device = await find_device(args.name, args.address, args.scan_timeout)
|
|
if device is None:
|
|
print(f"Device not found: {args.name}", flush=True)
|
|
if args.once:
|
|
return 2
|
|
await asyncio.sleep(args.retry_seconds)
|
|
continue
|
|
|
|
print(f"Connecting to {device.name or args.name} ({device.address})...", flush=True)
|
|
try:
|
|
async with BleakClient(device, timeout=args.connect_timeout) as client:
|
|
print("BLE connected. Streaming metrics.", flush=True)
|
|
while True:
|
|
line = sampler.sample_line()
|
|
await client.write_gatt_char(args.char_uuid, line.encode("ascii"), response=True)
|
|
print(line, flush=True)
|
|
if args.once:
|
|
return 0
|
|
await asyncio.sleep(args.interval)
|
|
except KeyboardInterrupt:
|
|
return 130
|
|
except Exception as exc:
|
|
print(f"BLE push failed: {exc}", file=sys.stderr, flush=True)
|
|
if args.once:
|
|
return 3
|
|
await asyncio.sleep(args.retry_seconds)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Push PC metrics to ESP32 RLCD over BLE.")
|
|
parser.add_argument("--name", default=DEFAULT_DEVICE_NAME)
|
|
parser.add_argument("--address", default=None)
|
|
parser.add_argument("--char-uuid", default=DEFAULT_CHAR_UUID)
|
|
parser.add_argument("--interval", type=float, default=1.0)
|
|
parser.add_argument("--scan-timeout", type=float, default=20.0)
|
|
parser.add_argument("--connect-timeout", type=float, default=20.0)
|
|
parser.add_argument("--retry-seconds", type=float, default=5.0)
|
|
parser.add_argument("--disk", default="C:\\")
|
|
parser.add_argument("--once", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
try:
|
|
return asyncio.run(push_loop(args))
|
|
except KeyboardInterrupt:
|
|
return 130
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|