ocdway/scripts/ocd-waybar
ocpway c4647dda6b Fix hang on click and reconnect after daemon restart
Two root causes shared one symptom: the waybar module froze after a
click and needed a waybar restart to recover.

1. oc-waybar-ctl used socat's UNIX-SENDTO, a bidirectional relay that
   can block on a datagram socket. waybar runs the click command on
   its per-module event thread and stops updating the module until the
   command returns, so a hung socat froze the module. Replace the
   socat/nc chain with a bounded python3 sendto (2s timeout). python3
   is guaranteed present since the daemon needs it.

2. The daemon ran oc raw json with a blocking subprocess.run on the
   main loop, so a slow network fetch (up to 60s) stalled the 5s
   rotation and queued clicks. Move fetch_ocd to a worker thread; the
   main loop adopts the result when it lands and never blocks. Rotation
   and click draining stay live during a fetch.

Reconnect after a daemon restart now works because the click helper no
longer wedges waybar's event thread; the module's interval-1 cat
re-reads state.json as soon as the new daemon rewrites it (verified:
state.json returns within 1s).
2026-07-30 04:43:45 +03:00

437 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 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)}"
lines = [head, ""]
for i, acct in enumerate(accounts):
marker = "\u25b8 " if i == (index % len(accounts)) else " "
# The lock emoji belongs only to the currently displayed row.
row_locked = locked and i == (index % len(accounts))
body = render_account_line(acct, now, row_locked, markup=TOOLTIP_MARKUP, with_name=True)
lines.append(f"{marker}{body}")
return "\n".join(lines)
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()
time.sleep(d.event_deadline())
d.stop()
if __name__ == "__main__":
main()