a9254fc70e
The main loop slept on time.sleep(event_deadline()), so a click sat in the datagram socket until the next 0.5s timer tick woke the loop. That added up to ~450ms of latency before the click affected state.json, on top of waybar's 1s re-poll. Replace the sleep with select() on the socket: select returns the moment a datagram arrives, so the loop drains and writes the click within milliseconds (measured 29ms) while still bounding idle sleep by the next timer deadline. Rotation cadence and click-during-fetch behavior are unchanged. The waybar custom-module interval stays 1: fractional intervals below 1s break custom modules (Waybar #4894 turns 0.5 into a 1ms busy loop), so the 0-1s re-poll is the floor and not safely tunable.
470 lines
15 KiB
Python
Executable File
470 lines
15 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# ocd-waybar: waybar daemon for the ocd opencode-limits dashboard.
|
|
#
|
|
# It polls `ocd raw json` every 30 s. It does not pass --force-refresh or
|
|
# --no-cache, so ocd's own cache and TTL stay in charge of politeness.
|
|
# It holds the 5 s rotation timer and the 1 h lock, renders one Pango line
|
|
# for the displayed account and a multi-line Pango tooltip for all accounts,
|
|
# and writes the result as JSON for the waybar custom/ocd module.
|
|
#
|
|
# Clicks arrive over a unix socket from ocd-waybar-ctl:
|
|
# advance -> reset the 5 s timer and move to the next account
|
|
# lock -> freeze the current account for one hour
|
|
# poll -> force an ocd poll now (still respects ocd's cache/TTL)
|
|
# quit -> stop the daemon
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import select
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from html import escape
|
|
from pathlib import Path
|
|
|
|
# --- tunables ---------------------------------------------------------------
|
|
ROTATE_SECS = 5 # normal auto-advance cadence
|
|
LOCK_SECS = 3600 # one hour lock
|
|
POLL_SECS = 30 # ocd re-poll cadence (cache serves most calls)
|
|
TICK_SECS = 0.5 # main loop resolution
|
|
BAR_WIDTH = 7 # nbsp cells per bar (~50 px at normal bar font)
|
|
STATE_DIR = Path.home() / ".local" / "state" / "ocd-waybar"
|
|
STATE_FILE = STATE_DIR / "state.json"
|
|
SOCK_FILE = Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")) / "ocd-waybar.sock"
|
|
OCD_BIN = os.environ.get("OCD_BIN", "ocd")
|
|
# Default False: waybar 0.15.0 renders the custom-module tooltip as plain
|
|
# text, so tags would show literally. Set True if your waybar supports
|
|
# Pango markup in custom tooltips (see Waybar PR #5034).
|
|
TOOLTIP_MARKUP = False
|
|
|
|
# zone colors
|
|
C_OK = "#43c267"
|
|
C_WARN = "#e0a030"
|
|
C_CRIT = "#e0533b"
|
|
C_LOCK = "#5aa6e0"
|
|
C_TRACK = "#39424d"
|
|
|
|
ZONE = {("ok",): 0, ("warn",): 1, ("crit",): 2}
|
|
# Each label is a template: {} is filled with the used value plus its
|
|
# unit, e.g. 88 and "%" -> "[88%/5h]". With no data the unit alone shows.
|
|
FRAME_LABELS = (("usage_5h", "[{}/5h]"), ("usage_week", "[{}/1w]"), ("usage_month", "[{}/1m]"))
|
|
|
|
|
|
def zone_color(pct, status=None, locked=False):
|
|
if locked:
|
|
return C_LOCK
|
|
if status == "rate-limited" or pct >= 90:
|
|
return C_CRIT
|
|
if pct >= 70:
|
|
return C_WARN
|
|
return C_OK
|
|
|
|
|
|
def zone_name(pct, status=None):
|
|
if status == "rate-limited" or pct >= 90:
|
|
return "ocd-crit"
|
|
if pct >= 70:
|
|
return "ocd-warn"
|
|
return "ocd-ok"
|
|
|
|
|
|
def parse_dt(s):
|
|
"""Parse an ISO 8601 string with possible nanosecond precision."""
|
|
if not s:
|
|
return None
|
|
s = s.strip()
|
|
if s.endswith("Z"):
|
|
s = s[:-1] + "+00:00"
|
|
# Python only handles up to 6 fractional digits. Trim the rest.
|
|
m = re.search(r"(\.\d{6})\d+", s)
|
|
if m:
|
|
s = s[:m.start()] + m.group(1) + s[m.end():]
|
|
try:
|
|
dt = datetime.fromisoformat(s)
|
|
except ValueError:
|
|
return None
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return dt
|
|
|
|
|
|
def rel_reset(resets_at, now):
|
|
"""Human string for time until reset, e.g. '4:39', '3d 1h', '20d'."""
|
|
if not resets_at:
|
|
return "\u2014"
|
|
dt = parse_dt(resets_at)
|
|
if dt is None:
|
|
return "\u2014"
|
|
sec = (dt - now).total_seconds()
|
|
if sec <= 0:
|
|
return "now"
|
|
days = int(sec // 86400)
|
|
rem = sec - days * 86400
|
|
hours = int(rem // 3600)
|
|
mins = int((rem % 3600) // 60)
|
|
if days >= 7:
|
|
return f"{days}d"
|
|
if days >= 1:
|
|
h = hours if hours else 0
|
|
return f"{days}d {h}h"
|
|
return f"{hours}:{mins:02d}"
|
|
|
|
|
|
def rel_ago(fetched_at, now):
|
|
if not fetched_at:
|
|
return "\u2014"
|
|
dt = parse_dt(fetched_at)
|
|
if dt is None:
|
|
return "\u2014"
|
|
sec = (now - dt).total_seconds()
|
|
if sec < 0:
|
|
sec = 0
|
|
if sec < 60:
|
|
return f"{int(sec)}s ago"
|
|
m = int(sec // 60)
|
|
if m < 60:
|
|
return f"{max(1, m)}m ago"
|
|
h = int(sec // 3600)
|
|
if h < 24:
|
|
return f"{h}h ago"
|
|
return f"{h // 24}d ago"
|
|
|
|
|
|
def bar_html(pct, status=None, locked=False):
|
|
pct = max(0.0, min(100.0, float(pct)))
|
|
fill = int(round(pct / 100.0 * BAR_WIDTH))
|
|
fill = max(0, min(BAR_WIDTH, fill))
|
|
empty = BAR_WIDTH - fill
|
|
fc = zone_color(pct, status, locked)
|
|
out = f'<span background="{fc}">{"\u00a0" * fill}</span>'
|
|
if empty:
|
|
out += f'<span background="{C_TRACK}">{"\u00a0" * empty}</span>'
|
|
return out
|
|
|
|
|
|
def bar_plain(pct, status=None, locked=False):
|
|
"""Plain block-character bar for tooltip text that may not parse markup."""
|
|
pct = max(0.0, min(100.0, float(pct)))
|
|
units = pct / 100.0 * BAR_WIDTH # cells, fractional
|
|
full = int(units)
|
|
frac = units - full
|
|
blocks = "▏▎▍▌▋▊▉" # one-eighth steps
|
|
out = "█" * full
|
|
if full < BAR_WIDTH and frac > 0:
|
|
idx = int(round(frac * 7)) - 1 # 7 fractional steps
|
|
idx = max(0, min(6, idx))
|
|
out += blocks[idx]
|
|
full += 1
|
|
out += "░" * (BAR_WIDTH - full) # light shade for the empty track
|
|
if locked and full < BAR_WIDTH:
|
|
out = out[:full] + "▓" + out[full + 1:] # dim mark on locked rows
|
|
return out[:BAR_WIDTH]
|
|
|
|
|
|
def frame_get(limits, key):
|
|
if not limits:
|
|
return None
|
|
return limits.get(key)
|
|
|
|
|
|
def pct_of(frame):
|
|
if frame is None:
|
|
return 0.0
|
|
total = float(frame.get("total") or 0) or 100.0
|
|
used = float(frame.get("used") or 0)
|
|
return used / total * 100.0 if total else 0.0
|
|
|
|
|
|
def worst_zone(account, locked=False):
|
|
if not account or not account.get("limits"):
|
|
return "ocd-ok"
|
|
worst = 0
|
|
names = ("usage_5h", "usage_week", "usage_month")
|
|
for k in names:
|
|
f = account["limits"].get(k)
|
|
if not f:
|
|
continue
|
|
zn = zone_name(pct_of(f), f.get("status"))
|
|
w = {"ocd-ok": 0, "ocd-warn": 1, "ocd-crit": 2}[zn]
|
|
worst = max(worst, w)
|
|
return ("ocd-ok", "ocd-warn", "ocd-crit")[worst]
|
|
|
|
|
|
def balance_str(account):
|
|
bal = account.get("balance") if account else None
|
|
if not bal:
|
|
return "\u2014"
|
|
remaining = float(bal.get("remaining") or 0.0)
|
|
unit = bal.get("unit") or ""
|
|
# Currency symbol prefixes; others append a space.
|
|
if unit and not unit.isalpha():
|
|
return f"{unit}{remaining:.2f}"
|
|
return f"{remaining:.2f} {unit}".strip()
|
|
|
|
|
|
def render_account_line(account, now, locked, markup=True, with_name=True):
|
|
"""Render one account: name + lock emoji (optional), then the three
|
|
bars with their percentage labels, then balance and update age."""
|
|
limits = account.get("limits") if account else None
|
|
parts = []
|
|
if with_name:
|
|
nm = account.get("account") if account else ""
|
|
if nm:
|
|
nm = escape(nm) if markup else nm
|
|
parts.append(("\U0001f512 " if locked else "") + nm)
|
|
segs = []
|
|
for key, label in FRAME_LABELS:
|
|
f = frame_get(limits, key)
|
|
if f is None:
|
|
filled = label.format("\u2014")
|
|
bar = "" if markup else ""
|
|
rel = "\u2014"
|
|
else:
|
|
used = f.get("used") or 0.0
|
|
unit = f.get("unit") or ""
|
|
val = f"{int(round(used))}{unit}"
|
|
filled = label.format(val)
|
|
bar = bar_html(pct_of(f), f.get("status"), locked) if markup else bar_plain(pct_of(f), f.get("status"), locked)
|
|
rel = rel_reset(f.get("resets_at"), now)
|
|
segs.append(f"{filled} {bar} {rel}")
|
|
parts.append(" ".join(segs))
|
|
bal = balance_str(account or {})
|
|
fa = rel_ago(account.get("fetched_at") if account else None, now)
|
|
parts.append(f"{bal} / {fa}")
|
|
return " ".join(parts)
|
|
|
|
|
|
def render_text(report, index, locked, now):
|
|
accounts = (report or {}).get("accounts") or []
|
|
if not accounts:
|
|
return ("ocd: no data", worst_zone(None))
|
|
idx = index % len(accounts)
|
|
acct = accounts[idx]
|
|
line = render_account_line(acct, now, locked, markup=True, with_name=True)
|
|
return (line, worst_zone(acct, locked))
|
|
|
|
|
|
def render_tooltip(report, index, locked, now):
|
|
accounts = (report or {}).get("accounts") or []
|
|
if not accounts:
|
|
return "ocd: waiting for first poll"
|
|
gen = parse_dt((report or {}).get("generated_at"))
|
|
head = "opencode limits"
|
|
if gen:
|
|
head += f" \u00b7 updated {rel_ago(report.get('generated_at'), now)}"
|
|
# Each window (5h/1w/1m) gets its own row; a blank row separates
|
|
# accounts for a bigger gap. Waybar custom-module tooltips need \r
|
|
# for line breaks (\n makes the tooltip vanish; see Waybar #3153).
|
|
windows = (("usage_5h", "5h"), ("usage_week", "1w"), ("usage_month", "1m"))
|
|
blocks = [head, ""]
|
|
cur = index % len(accounts)
|
|
for i, acct in enumerate(accounts):
|
|
is_cur = i == cur
|
|
marker = "\u25b8 " if is_cur else " "
|
|
row_locked = locked and is_cur
|
|
nm = (acct or {}).get("account") or ""
|
|
lock = " \U0001f512" if row_locked else ""
|
|
bal = balance_str(acct or {})
|
|
age = rel_ago((acct or {}).get("fetched_at"), now)
|
|
blocks.append(f"{marker}{nm}{lock} {bal} \u00b7 {age}")
|
|
limits = (acct or {}).get("limits") or {}
|
|
for key, wl in windows:
|
|
f = frame_get(limits, key)
|
|
if f is None:
|
|
pct = 0.0
|
|
status = None
|
|
reset = "\u2014"
|
|
else:
|
|
pct = pct_of(f)
|
|
status = f.get("status")
|
|
reset = rel_reset(f.get("resets_at"), now)
|
|
bar = bar_plain(pct, status, row_locked)
|
|
blocks.append(f" {wl} [{int(round(pct)):>3}%] {bar} {reset}")
|
|
blocks.append("")
|
|
return "\r".join(blocks)
|
|
|
|
|
|
def write_state(text, tooltip, klass):
|
|
STATE_DIR.mkdir(parents=True, exist_ok=True)
|
|
payload = {
|
|
"text": text,
|
|
"tooltip": tooltip,
|
|
"class": klass,
|
|
"alt": "ocd",
|
|
}
|
|
tmp = STATE_FILE.with_suffix(".json.tmp")
|
|
tmp.write_text(json.dumps(payload, ensure_ascii=False))
|
|
tmp.replace(STATE_FILE)
|
|
|
|
|
|
def fetch_ocd():
|
|
"""Run `ocd raw json`. Returns parsed dict or None on failure."""
|
|
try:
|
|
proc = subprocess.run(
|
|
[OCD_BIN, "raw", "json"],
|
|
capture_output=True, text=True, timeout=60,
|
|
)
|
|
if proc.returncode != 0:
|
|
return None
|
|
return json.loads(proc.stdout)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
class Daemon:
|
|
def __init__(self):
|
|
self.report = None
|
|
self.index = 0
|
|
self.last_advance = time.monotonic()
|
|
self.locked_until = 0.0
|
|
self.next_poll = 0.0
|
|
self.last_text = None
|
|
self.last_key = None
|
|
self.sock = None
|
|
self.running = True
|
|
# The ocd poll runs in a worker thread so the main loop never
|
|
# blocks on a slow network fetch. Rotation and clicks stay live.
|
|
self._fetch_thread = None
|
|
self._fresh_report = None
|
|
self._fetch_lock = threading.Lock()
|
|
|
|
def event_deadline(self):
|
|
now = time.monotonic()
|
|
deadlines = [now + TICK_SECS, self.next_poll]
|
|
if not self.is_locked(now):
|
|
deadlines.append(self.last_advance + ROTATE_SECS)
|
|
else:
|
|
deadlines.append(self.locked_until)
|
|
return max(0.0, min(deadlines) - now)
|
|
|
|
def is_locked(self, now=None):
|
|
if now is None:
|
|
now = time.monotonic()
|
|
return self.locked_until and now < self.locked_until
|
|
|
|
def advance(self):
|
|
n = len((self.report or {}).get("accounts") or [])
|
|
if n:
|
|
self.index = (self.index + 1) % n
|
|
self.last_advance = time.monotonic()
|
|
|
|
def lock(self):
|
|
self.locked_until = time.monotonic() + LOCK_SECS
|
|
|
|
def _fetch_worker(self):
|
|
rep = fetch_ocd()
|
|
if rep:
|
|
with self._fetch_lock:
|
|
self._fresh_report = rep
|
|
|
|
def maybe_tick(self):
|
|
now = time.monotonic()
|
|
# Kick off a poll on a worker thread; never block the main loop.
|
|
if now >= self.next_poll and (self._fetch_thread is None or not self._fetch_thread.is_alive()):
|
|
self.next_poll = now + POLL_SECS
|
|
self._fetch_thread = threading.Thread(target=self._fetch_worker, daemon=True)
|
|
self._fetch_thread.start()
|
|
# Adopt a finished poll's result if one is ready.
|
|
if self._fetch_thread is not None and not self._fetch_thread.is_alive():
|
|
self._fetch_thread = None
|
|
with self._fetch_lock:
|
|
rep = self._fresh_report
|
|
self._fresh_report = None
|
|
if rep:
|
|
self.report = rep
|
|
n = len(rep.get("accounts") or [])
|
|
if n and self.index >= n:
|
|
self.index %= n
|
|
if not self.is_locked(now) and (now - self.last_advance) >= ROTATE_SECS:
|
|
self.advance()
|
|
if self.is_locked(now) and now >= self.locked_until:
|
|
self.locked_until = 0.0
|
|
|
|
def render_and_write(self):
|
|
now = datetime.now(timezone.utc)
|
|
text, zn = render_text(self.report, self.index, self.is_locked(), now)
|
|
tooltip = render_tooltip(self.report, self.index, self.is_locked(), now)
|
|
klass = ("ocd " + zn + (" ocd-locked" if self.is_locked() else "")).strip()
|
|
# Write only when something changed. Displayed relative times
|
|
# are minute-resolution, so the file updates a few times per minute.
|
|
key = (text, tooltip, klass)
|
|
if key != self.last_key:
|
|
write_state(text, tooltip, klass)
|
|
self.last_key = key
|
|
|
|
def handle_cmd(self, cmd):
|
|
cmd = (cmd or "").strip()
|
|
if cmd == "advance":
|
|
self.advance()
|
|
self.locked_until = 0.0
|
|
elif cmd == "lock":
|
|
self.lock()
|
|
elif cmd == "poll":
|
|
self.next_poll = 0.0
|
|
elif cmd == "quit":
|
|
self.running = False
|
|
|
|
def setup_socket(self):
|
|
try:
|
|
os.unlink(SOCK_FILE)
|
|
except FileNotFoundError:
|
|
pass
|
|
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
|
|
self.sock.bind(str(SOCK_FILE))
|
|
self.sock.setblocking(False)
|
|
os.chmod(SOCK_FILE, 0o600)
|
|
|
|
def drain_socket(self):
|
|
if self.sock is None:
|
|
return
|
|
while True:
|
|
try:
|
|
data, _ = self.sock.recvfrom(4096)
|
|
except BlockingIOError:
|
|
return
|
|
self.handle_cmd(data.decode("utf-8", "replace"))
|
|
|
|
def stop(self, *_):
|
|
self.running = False
|
|
try:
|
|
if self.sock:
|
|
self.sock.close()
|
|
os.unlink(SOCK_FILE)
|
|
except FileNotFoundError:
|
|
pass
|
|
sys.exit(0)
|
|
|
|
|
|
def main():
|
|
STATE_DIR.mkdir(parents=True, exist_ok=True)
|
|
d = Daemon()
|
|
d.setup_socket()
|
|
signal.signal(signal.SIGTERM, d.stop)
|
|
signal.signal(signal.SIGINT, d.stop)
|
|
# force first poll
|
|
d.next_poll = 0.0
|
|
while d.running:
|
|
d.maybe_tick()
|
|
d.drain_socket()
|
|
d.render_and_write()
|
|
# Sleep until the next tick OR a click, whichever is first.
|
|
# select on the socket lets a click wake the loop at once instead
|
|
# of waiting up to TICK_SECS for the next timer tick.
|
|
timeout = d.event_deadline()
|
|
if d.sock is not None:
|
|
try:
|
|
select.select([d.sock], [], [], timeout)
|
|
except (OSError, ValueError):
|
|
time.sleep(timeout)
|
|
else:
|
|
time.sleep(timeout)
|
|
d.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |