In a managed Agent-SDK environment `claude login` is unavailable, so once
the stored access token expired the 5h/7d rate-limit % went dark forever.
The daemon now renews its own access token from the stored refreshToken via
the standard Claude Code OAuth refresh grant and writes the rotated tokens
back atomically (temp + os.replace, so a crash can't corrupt the file that
both the daemon and Claude Code read).
- proactive: refresh ~300s before expiry, checked each poll cycle
- reactive: on a genuine API 401, force one refresh + retry the poll once
- refresh tokens rotate on every use -> always persisted
- never raises: any refresh failure is logged and the loop continues
- secrets never logged (only lengths / expiry); error bodies are {"error":..}
Endpoint + client_id verified empirically against a live refresh
(HTTP 200, expires_in=28800).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
801 lines
33 KiB
Python
801 lines
33 KiB
Python
#!/usr/bin/env python3
|
||
"""Claude Usage Tracker Daemon — Windows (Phase 2).
|
||
|
||
Reads the Claude OAuth token from the native-Windows credentials path and
|
||
polls the Anthropic API for rate-limit utilization data. BLE glue added in
|
||
later plans.
|
||
"""
|
||
|
||
import asyncio
|
||
import datetime
|
||
import json
|
||
import logging
|
||
import logging.handlers
|
||
import os
|
||
import re
|
||
import signal
|
||
import sys
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
from bleak import BleakClient, BleakScanner
|
||
from bleak.exc import BleakError
|
||
|
||
DEVICE_NAME = "Clawdmeter" # must match firmware ble.cpp DEVICE_NAME
|
||
# Optional direct-connect address. Once the device is bonded as a BLE HID
|
||
# keyboard, Windows keeps it connected, so it stops advertising and an active
|
||
# scan (find_device_by_name) never sees it. Setting CLAWDMETER_ADDRESS lets the
|
||
# daemon connect straight to the bonded device by address via WinRT.
|
||
DEVICE_ADDRESS = os.environ.get("CLAWDMETER_ADDRESS")
|
||
SERVICE_UUID = "4c41555a-4465-7669-6365-000000000001"
|
||
RX_CHAR_UUID = "4c41555a-4465-7669-6365-000000000002"
|
||
REQ_CHAR_UUID = "4c41555a-4465-7669-6365-000000000004"
|
||
|
||
POLL_INTERVAL = 60
|
||
TICK = 5
|
||
SCAN_TIMEOUT = 8.0
|
||
CONNECT_RETRIES = 3 # D-01: attempts before giving up on a device
|
||
CONNECT_RETRY_DELAY = 2.0 # D-01: seconds between failed connect attempts
|
||
ZOMBIE_BREAK_LIMIT = 1 # D-03: consecutive write failures before abandoning a half-open link
|
||
# N=1: breaks at T=60s, leaves ~60s headroom for reconnect+poll inside 120s SLA
|
||
# N=2 would bust the 120s budget before reconnect even begins
|
||
RECONNECT_BACKOFF_CAP = 8 # D-05: fast-reconnect cap (seconds); keeps stacked retries inside 120s SLA
|
||
# ~5–10s band per CONTEXT.md Claude's Discretion; 8 chosen as middle ground
|
||
|
||
API_URL = "https://api.anthropic.com/v1/messages"
|
||
API_HEADERS_TEMPLATE = {
|
||
"anthropic-version": "2023-06-01",
|
||
"anthropic-beta": "oauth-2025-04-20",
|
||
"Content-Type": "application/json",
|
||
"User-Agent": "claude-code/2.1.5",
|
||
}
|
||
API_BODY = {
|
||
"model": "claude-haiku-4-5-20251001",
|
||
"max_tokens": 1,
|
||
"messages": [{"role": "user", "content": "hi"}],
|
||
}
|
||
|
||
# --- OAuth self-refresh -------------------------------------------------------
|
||
# In a managed (Agent SDK) environment `claude login` is unavailable, so once the
|
||
# stored access token expires nothing renews it and the 5h/7d rate-limit % goes
|
||
# dark permanently. The credentials file ships a refreshToken, so the daemon
|
||
# renews its own access token via the standard Claude Code OAuth refresh grant
|
||
# and writes the (rotated) tokens back. Endpoint + client_id verified empirically
|
||
# against a live refresh: HTTP 200, refresh_token rotates, expires_in=28800 (8h).
|
||
OAUTH_TOKEN_URL = "https://console.anthropic.com/v1/oauth/token"
|
||
OAUTH_CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
|
||
TOKEN_REFRESH_SKEW = 300 # refresh this many seconds BEFORE expiry (proactive)
|
||
|
||
# Anthropic list prices, USD per million tokens, keyed by a substring of the
|
||
# model id. Used to show the "equivalent API cost" of today's Claude Code usage
|
||
# — subscription users don't actually pay this; it's the ccusage-style flex.
|
||
# Verified against ccusage/LiteLLM: cost reproduces to 100%. The rates follow
|
||
# Anthropic's fixed structure (output 5x, cache-write 2x, cache-read 0.1x of the
|
||
# base input rate). Opus 4.x is $5/$25 — NOT the older $15/$75.
|
||
PRICING = {
|
||
"opus": {"in": 5.0, "out": 25.0, "cache_w": 10.0, "cache_r": 0.50},
|
||
"sonnet": {"in": 3.0, "out": 15.0, "cache_w": 6.0, "cache_r": 0.30},
|
||
"haiku": {"in": 0.75, "out": 3.75, "cache_w": 1.50, "cache_r": 0.075},
|
||
}
|
||
_DEFAULT_PRICE = PRICING["opus"]
|
||
|
||
|
||
def _price_for(model: str) -> dict:
|
||
m = (model or "").lower()
|
||
for key, price in PRICING.items():
|
||
if key in m:
|
||
return price
|
||
return _DEFAULT_PRICE
|
||
|
||
|
||
def compute_today_usage() -> dict:
|
||
"""Sum today's Claude Code token usage + equivalent API cost across all local
|
||
project transcripts (~/.claude/projects/**/*.jsonl).
|
||
|
||
Reads only files modified today (a session that touched today has mtime
|
||
today), so the scan stays cheap even with a large transcript history. Each
|
||
assistant line carries message.usage (input/output/cache token counts) and a
|
||
UTC timestamp; per-line timestamps gate to today so a session spanning
|
||
midnight is split correctly. Assistant turns are de-duplicated by
|
||
(messageId, requestId): Claude Code copies history into new transcript files
|
||
on compaction/resume, so the same API turn appears in several files and a
|
||
naive sum over-counts (it inflated the total ~2.7x). Matches ccusage.
|
||
Returns the compact BLE fields tk/tc/tn/to.
|
||
"""
|
||
base = Path.home() / ".claude" / "projects"
|
||
today = datetime.date.today()
|
||
midnight = datetime.datetime.combine(today, datetime.time.min)
|
||
total_tokens = 0
|
||
output_tokens = 0
|
||
cost = 0.0
|
||
messages = 0
|
||
seen = set()
|
||
try:
|
||
files = list(base.glob("**/*.jsonl"))
|
||
except OSError:
|
||
return {"tk": 0, "tc": 0, "tn": 0}
|
||
for f in files:
|
||
try:
|
||
if datetime.datetime.fromtimestamp(f.stat().st_mtime) < midnight:
|
||
continue
|
||
except OSError:
|
||
continue
|
||
try:
|
||
with open(f, encoding="utf-8") as fh:
|
||
for line in fh:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
obj = json.loads(line)
|
||
except (json.JSONDecodeError, ValueError):
|
||
continue
|
||
msg = obj.get("message") if isinstance(obj, dict) else None
|
||
if not isinstance(msg, dict):
|
||
continue
|
||
usage = msg.get("usage")
|
||
if not isinstance(usage, dict):
|
||
continue
|
||
ts = obj.get("timestamp")
|
||
if not ts:
|
||
continue
|
||
try:
|
||
d = datetime.datetime.fromisoformat(
|
||
ts.replace("Z", "+00:00")).astimezone().date()
|
||
except ValueError:
|
||
continue
|
||
if d != today:
|
||
continue
|
||
# De-dup the same API turn copied across resumed/compacted files.
|
||
mid = obj.get("messageId") or msg.get("id")
|
||
rid = obj.get("requestId")
|
||
if mid is not None and rid is not None:
|
||
key = (mid, rid)
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
inp = usage.get("input_tokens", 0) or 0
|
||
out = usage.get("output_tokens", 0) or 0
|
||
cw = usage.get("cache_creation_input_tokens", 0) or 0
|
||
cr = usage.get("cache_read_input_tokens", 0) or 0
|
||
total_tokens += inp + out + cw + cr
|
||
output_tokens += out
|
||
messages += 1
|
||
p = _price_for(msg.get("model"))
|
||
cost += (inp * p["in"] + out * p["out"]
|
||
+ cw * p["cache_w"] + cr * p["cache_r"]) / 1_000_000
|
||
except OSError:
|
||
continue
|
||
return {"tk": int(total_tokens), "tc": int(round(cost * 100)),
|
||
"tn": messages, "to": int(output_tokens)}
|
||
|
||
|
||
def _build_file_logger() -> logging.Logger | None:
|
||
"""Create a rotating file logger for field diagnostics, or None.
|
||
|
||
Autostart launches the tray under pythonw.exe, which has no console — stdout
|
||
is discarded (and is in fact None, making print() unsafe). A rotating file is
|
||
then the ONLY trail when the daemon stalls in the field. Windows-only: on the
|
||
Linux dev box / CI the console print() suffices, and gating to win32 keeps the
|
||
pure-helper unit tests from writing stray log files.
|
||
"""
|
||
if sys.platform != "win32":
|
||
return None
|
||
logger = logging.getLogger("clawdmeter.daemon")
|
||
if logger.handlers:
|
||
return logger # idempotent across re-import (tray imports this module)
|
||
base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local"))
|
||
path = base / "Clawdmeter" / "daemon.log"
|
||
try:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
handler = logging.handlers.RotatingFileHandler(
|
||
path, maxBytes=512 * 1024, backupCount=3, encoding="utf-8"
|
||
)
|
||
except OSError:
|
||
return None # best-effort — logging setup must never stop the daemon
|
||
handler.setFormatter(logging.Formatter("%(asctime)s %(message)s", "%Y-%m-%d %H:%M:%S"))
|
||
logger.addHandler(handler)
|
||
logger.setLevel(logging.INFO)
|
||
logger.propagate = False
|
||
return logger
|
||
|
||
|
||
_FILE_LOGGER = _build_file_logger()
|
||
|
||
|
||
def log(msg: str) -> None:
|
||
line = f"[{time.strftime('%H:%M:%S')}] {msg}"
|
||
# Under pythonw sys.stdout is None and print() would raise — guard it so a
|
||
# missing console can never crash the daemon thread (the silent-freeze mode).
|
||
try:
|
||
print(line, flush=True)
|
||
except (OSError, ValueError, AttributeError, RuntimeError):
|
||
pass
|
||
if _FILE_LOGGER is not None:
|
||
_FILE_LOGGER.info(msg)
|
||
|
||
|
||
class AuthError(Exception):
|
||
"""Raised by poll_api on a genuine 401/403 — the token really is expired or
|
||
invalid and the user must re-run `claude login`. Distinct from a None return,
|
||
which means a TRANSIENT failure (network/DNS, timeout, rate-limit, 5xx) that
|
||
must NOT be mislabeled as a token problem (SC#5: a boot-time `getaddrinfo
|
||
failed` DNS blip wrongly fired the 'token expired' toast)."""
|
||
|
||
|
||
async def poll_api(token: str) -> dict | None:
|
||
headers = dict(API_HEADERS_TEMPLATE)
|
||
headers["Authorization"] = f"Bearer {token}"
|
||
try:
|
||
async with httpx.AsyncClient(timeout=20.0) as http:
|
||
resp = await http.post(API_URL, headers=headers, json=API_BODY)
|
||
except httpx.HTTPError as e:
|
||
# Network/DNS/timeout — transient. Return None (no toast), retry next tick.
|
||
log(f"API call failed: {e}")
|
||
return None
|
||
if resp.status_code in (401, 403):
|
||
# Genuine auth rejection — the ONLY case that warrants the actionable
|
||
# "run claude login" toast.
|
||
log(f"API HTTP {resp.status_code}: {resp.text[:200]}")
|
||
raise AuthError(resp.status_code)
|
||
if resp.status_code >= 400:
|
||
# Other 4xx/5xx (rate-limit, server error) — transient, not a token issue.
|
||
log(f"API HTTP {resp.status_code}: {resp.text[:200]}")
|
||
return None
|
||
|
||
def hdr(name: str, default: str = "0") -> str:
|
||
return resp.headers.get(name, default)
|
||
|
||
now = time.time()
|
||
|
||
def reset_minutes(reset_ts: str) -> int:
|
||
try:
|
||
r = float(reset_ts)
|
||
except ValueError:
|
||
return 0
|
||
mins = (r - now) / 60.0
|
||
return int(round(mins)) if mins > 0 else 0
|
||
|
||
def pct(util: str) -> int:
|
||
try:
|
||
return int(round(float(util) * 100))
|
||
except ValueError:
|
||
return 0
|
||
|
||
# Rate-limit utilization only. The local token usage and the "ok" flag are
|
||
# merged by the caller (connect_and_run) so the Session screen keeps updating
|
||
# even when this call fails (expired token / API down) — see decoupling there.
|
||
return {
|
||
"s": pct(hdr("anthropic-ratelimit-unified-5h-utilization")),
|
||
"sr": reset_minutes(hdr("anthropic-ratelimit-unified-5h-reset")),
|
||
"w": pct(hdr("anthropic-ratelimit-unified-7d-utilization")),
|
||
"wr": reset_minutes(hdr("anthropic-ratelimit-unified-7d-reset")),
|
||
"st": hdr("anthropic-ratelimit-unified-5h-status", "unknown"),
|
||
}
|
||
|
||
|
||
async def scan_for_device():
|
||
"""Scan for DEVICE_NAME and return the BLEDevice, or None."""
|
||
log(f"Scanning for '{DEVICE_NAME}' ({SCAN_TIMEOUT}s)...")
|
||
device = await BleakScanner.find_device_by_name(DEVICE_NAME, timeout=SCAN_TIMEOUT)
|
||
if device:
|
||
log(f"Found: {device.address}")
|
||
return device # BLEDevice
|
||
if DEVICE_ADDRESS:
|
||
# Bonded device that's connected to Windows won't advertise — connect by address.
|
||
log(f"Not advertising; connecting by bonded address {DEVICE_ADDRESS}")
|
||
return DEVICE_ADDRESS # address string — BleakClient accepts it on WinRT
|
||
return device # None
|
||
|
||
|
||
class Session:
|
||
def __init__(self, client: BleakClient) -> None:
|
||
self.client = client
|
||
self.refresh_requested = asyncio.Event()
|
||
|
||
def _on_refresh(self, _char, _data: bytearray) -> None:
|
||
log("Refresh requested by device")
|
||
self.refresh_requested.set()
|
||
|
||
async def setup_refresh_subscription(self) -> None:
|
||
# The refresh subscription is optional — the 60s poll loop works without it.
|
||
# WinRT's start_notify() CCCD write can raise a raw OSError/WinError (not
|
||
# wrapped as BleakError) when the peer GATT server is transiently unavailable,
|
||
# e.g. a just-power-cycled ESP32 whose server is not yet ready (G-03-01, SC#3).
|
||
# Degrade gracefully instead of crashing the daemon so it stays single-process
|
||
# across a power-cycle reconnect (SC#4, no restart).
|
||
try:
|
||
await self.client.start_notify(REQ_CHAR_UUID, self._on_refresh)
|
||
except (BleakError, ValueError, OSError) as e:
|
||
log(f"Refresh subscription unavailable: {e}")
|
||
|
||
async def write_payload(self, payload: dict) -> bool:
|
||
data = json.dumps(payload, separators=(",", ":")).encode()
|
||
log(f"Sending: {data.decode()}")
|
||
try:
|
||
await self.client.write_gatt_char(RX_CHAR_UUID, data, response=False)
|
||
return True
|
||
except (BleakError, OSError) as e:
|
||
# WinRT can raise a raw OSError/WinError (NOT wrapped as BleakError)
|
||
# when the peer GATT server goes transiently unavailable mid-write —
|
||
# the same failure class setup_refresh_subscription() guards against.
|
||
# Returning False trips the zombie-link break -> clean reconnect,
|
||
# rather than an uncaught exception killing the daemon thread (the
|
||
# silent-freeze failure mode, SC#2 field report).
|
||
log(f"Write failed: {e}")
|
||
return False
|
||
|
||
|
||
def _extract_access_token(blob: str) -> str | None:
|
||
"""Pull the accessToken out of a credentials blob.
|
||
|
||
Claude Code stores credentials as a JSON object; the blob may also be
|
||
nested ({"claudeAiOauth": {"accessToken": "..."}}). Fall back to a
|
||
regex match so unexpected shapes still work, and finally treat the
|
||
blob as a raw token if nothing else matches.
|
||
"""
|
||
blob = blob.strip()
|
||
if not blob:
|
||
return None
|
||
try:
|
||
data = json.loads(blob)
|
||
except json.JSONDecodeError:
|
||
data = None
|
||
if isinstance(data, dict):
|
||
# direct: {"accessToken": "..."}
|
||
tok = data.get("accessToken")
|
||
if isinstance(tok, str) and tok.strip():
|
||
return tok
|
||
# nested: {"claudeAiOauth": {"accessToken": "..."}}
|
||
for v in data.values():
|
||
if isinstance(v, dict):
|
||
tok = v.get("accessToken")
|
||
if isinstance(tok, str) and tok.strip():
|
||
return tok
|
||
m = re.search(r'"accessToken"\s*:\s*"([^"]+)"', blob)
|
||
if m:
|
||
return m.group(1)
|
||
# Raw token (no JSON wrapper) — must look plausible (sk-ant-... etc.)
|
||
if re.fullmatch(r"[A-Za-z0-9_\-.~+/=]{20,}", blob):
|
||
return blob
|
||
return None
|
||
|
||
|
||
def _windows_credential_candidates() -> list[Path]:
|
||
"""Return the ordered list of credential file paths to probe (first hit wins).
|
||
|
||
Priority:
|
||
1. CLAUDE_CREDENTIALS_PATH env override (D-03, project-specific)
|
||
2. CLAUDE_CONFIG_DIR env override (official Claude override)
|
||
3. D-02 candidate list: home/.claude, LOCALAPPDATA/Claude, APPDATA/Claude
|
||
"""
|
||
# Priority 1: project-specific env override (D-03)
|
||
if override := os.environ.get("CLAUDE_CREDENTIALS_PATH"):
|
||
return [Path(override)]
|
||
# Priority 2: official CLAUDE_CONFIG_DIR env override
|
||
if config_dir := os.environ.get("CLAUDE_CONFIG_DIR"):
|
||
return [Path(config_dir) / ".credentials.json"]
|
||
# Priority 3: D-02 candidate list — first hit wins
|
||
home = Path.home()
|
||
local_appdata = Path(os.environ.get("LOCALAPPDATA", home / "AppData" / "Local"))
|
||
appdata = Path(os.environ.get("APPDATA", home / "AppData" / "Roaming"))
|
||
return [
|
||
home / ".claude" / ".credentials.json", # primary (confirmed by docs)
|
||
local_appdata / "Claude" / ".credentials.json", # fallback 2
|
||
appdata / "Claude" / ".credentials.json", # fallback 3
|
||
]
|
||
|
||
|
||
def read_token() -> str | None:
|
||
"""Read the Claude OAuth access token from the first available credential file."""
|
||
for path in _windows_credential_candidates():
|
||
try:
|
||
return _extract_access_token(path.read_text(encoding="utf-8"))
|
||
except OSError:
|
||
continue
|
||
return None
|
||
|
||
|
||
def _read_expiry() -> str:
|
||
"""Return human-readable expiry from the first-hit credentials file.
|
||
|
||
Reads claudeAiOauth.expiresAt (epoch milliseconds — JS convention).
|
||
Divides by 1000 before passing to fromtimestamp (Python expects seconds).
|
||
Returns 'expiry unknown' on any parse failure.
|
||
"""
|
||
for path in _windows_credential_candidates():
|
||
try:
|
||
raw = path.read_text(encoding="utf-8")
|
||
except OSError:
|
||
continue
|
||
try:
|
||
data = json.loads(raw)
|
||
oauth = data.get("claudeAiOauth", {})
|
||
expires_ms = oauth.get("expiresAt")
|
||
if expires_ms is None:
|
||
return "expiry unknown"
|
||
# CRITICAL: expiresAt is JS-convention epoch milliseconds; divide by 1000
|
||
# before fromtimestamp (Python expects seconds). Raw value -> year ~57000.
|
||
dt = datetime.datetime.fromtimestamp(
|
||
expires_ms / 1000, tz=datetime.timezone.utc
|
||
)
|
||
return dt.strftime("%Y-%m-%d %H:%M UTC")
|
||
except (TypeError, ValueError, OSError, AttributeError, json.JSONDecodeError):
|
||
return "expiry unknown"
|
||
return "expiry unknown"
|
||
|
||
|
||
def _read_credentials() -> tuple[Path, dict] | None:
|
||
"""Return (path, parsed-JSON dict) for the first readable credentials file.
|
||
|
||
Same probe order as read_token() so the file we refresh is the file we read
|
||
the token from.
|
||
"""
|
||
for path in _windows_credential_candidates():
|
||
try:
|
||
return path, json.loads(path.read_text(encoding="utf-8"))
|
||
except (OSError, json.JSONDecodeError):
|
||
continue
|
||
return None
|
||
|
||
|
||
def _write_credentials_atomic(path: Path, data: dict) -> None:
|
||
"""Write credentials JSON via temp-file + os.replace (atomic same-volume).
|
||
|
||
Avoids leaving a half-written .credentials.json if the process dies mid-write
|
||
— a corrupt credentials file would break BOTH the daemon and Claude Code.
|
||
"""
|
||
tmp = path.with_name(path.name + ".tmp")
|
||
tmp.write_text(json.dumps(data, indent=2), encoding="utf-8")
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def _token_expired(oauth: dict, skew: int = TOKEN_REFRESH_SKEW) -> bool:
|
||
"""True if the access token is within `skew` seconds of expiry (or past it).
|
||
|
||
Unknown/absent expiresAt -> False: don't refresh blindly on every cycle; the
|
||
API-401 path (force=True) drives a refresh if the token is actually rejected.
|
||
"""
|
||
exp = oauth.get("expiresAt")
|
||
if not isinstance(exp, (int, float)):
|
||
return False
|
||
return (exp / 1000.0) <= (time.time() + skew)
|
||
|
||
|
||
async def refresh_token_if_needed(force: bool = False) -> bool:
|
||
"""Renew the OAuth access token from the stored refresh token when it is
|
||
expired/near-expiry (or force=True), writing the rotated tokens back atomically.
|
||
|
||
Returns True iff a refresh succeeded and credentials were updated. Never
|
||
raises — every failure is logged and returned as False so a refresh hiccup
|
||
can never take down the poll loop.
|
||
"""
|
||
rc = _read_credentials()
|
||
if rc is None:
|
||
return False
|
||
path, data = rc
|
||
oauth = data.get("claudeAiOauth")
|
||
if not isinstance(oauth, dict):
|
||
return False
|
||
refresh = oauth.get("refreshToken")
|
||
if not isinstance(refresh, str) or not refresh.strip():
|
||
return False
|
||
if not force and not _token_expired(oauth):
|
||
return False
|
||
|
||
body = {
|
||
"grant_type": "refresh_token",
|
||
"refresh_token": refresh,
|
||
"client_id": OAUTH_CLIENT_ID,
|
||
}
|
||
try:
|
||
async with httpx.AsyncClient(timeout=20.0) as http:
|
||
resp = await http.post(
|
||
OAUTH_TOKEN_URL, json=body,
|
||
headers={"Content-Type": "application/json"},
|
||
)
|
||
except httpx.HTTPError as e:
|
||
log(f"Token refresh network error: {e}")
|
||
return False
|
||
if resp.status_code != 200:
|
||
# Error bodies are {"error": ...} — no secrets; safe to log a snippet.
|
||
log(f"Token refresh HTTP {resp.status_code}: {resp.text[:200]}")
|
||
return False
|
||
try:
|
||
tok = resp.json()
|
||
except ValueError:
|
||
log("Token refresh: non-JSON response")
|
||
return False
|
||
|
||
new_access = tok.get("access_token")
|
||
if not new_access:
|
||
log("Token refresh: response had no access_token")
|
||
return False
|
||
oauth["accessToken"] = new_access
|
||
if tok.get("refresh_token"):
|
||
oauth["refreshToken"] = tok["refresh_token"] # refresh tokens rotate
|
||
if tok.get("expires_in"):
|
||
oauth["expiresAt"] = int((time.time() + float(tok["expires_in"])) * 1000)
|
||
try:
|
||
_write_credentials_atomic(path, data)
|
||
except OSError as e:
|
||
log(f"Token refresh: could not write credentials: {e}")
|
||
return False
|
||
log(f"OAuth token refreshed; new expiry {_read_expiry()}")
|
||
return True
|
||
|
||
|
||
async def _wait_first(*events: asyncio.Event, timeout: float) -> None:
|
||
"""Return when any of `events` is set, or after `timeout` seconds.
|
||
|
||
Lets the poll loop's TICK wait wake immediately on a stop signal (clean,
|
||
responsive Quit) without losing the refresh-request wakeup — instead of
|
||
waiting only on refresh_requested and re-checking stop_event up to TICK
|
||
later. Cancels and drains the loser tasks so they don't warn.
|
||
"""
|
||
tasks = [asyncio.ensure_future(e.wait()) for e in events]
|
||
try:
|
||
await asyncio.wait(tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED)
|
||
finally:
|
||
for t in tasks:
|
||
t.cancel()
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
|
||
async def connect_and_run(device, stop_event: asyncio.Event, tray_state=None) -> bool:
|
||
"""Connect to device and poll until disconnected or stopped.
|
||
|
||
Returns True if at least one successful write occurred.
|
||
"""
|
||
addr = device if isinstance(device, str) else device.address
|
||
log(f"Connecting to {addr}...")
|
||
# D-01: retry wrapper — defeats WinRT post-wake failure modes
|
||
# (Could not get GATT services: Unreachable, stale is_connected).
|
||
# Rebuild a fresh BleakClient each attempt (locked D-05 recipe).
|
||
client = None
|
||
for attempt in range(CONNECT_RETRIES):
|
||
# D-05: pass BLEDevice (not address string), address_type="random" (NimBLE
|
||
# static-random), use_cached_services=False (DIY firmware — WinRT GATT cache
|
||
# may be stale after firmware reflash).
|
||
client = BleakClient(
|
||
device,
|
||
address_type="random",
|
||
use_cached_services=False,
|
||
)
|
||
try:
|
||
await client.connect()
|
||
except (BleakError, asyncio.TimeoutError) as e:
|
||
log(f"Connection attempt {attempt + 1}/{CONNECT_RETRIES} failed: {e}")
|
||
try:
|
||
await client.disconnect()
|
||
except BleakError:
|
||
pass
|
||
if attempt < CONNECT_RETRIES - 1:
|
||
await asyncio.sleep(CONNECT_RETRY_DELAY)
|
||
continue
|
||
|
||
if not client.is_connected:
|
||
log(f"Connection attempt {attempt + 1}/{CONNECT_RETRIES} failed (not connected)")
|
||
try:
|
||
await client.disconnect()
|
||
except BleakError:
|
||
pass
|
||
if attempt < CONNECT_RETRIES - 1:
|
||
await asyncio.sleep(CONNECT_RETRY_DELAY)
|
||
continue
|
||
|
||
# Connected successfully
|
||
break
|
||
else:
|
||
log(f"Connection failed after {CONNECT_RETRIES} attempts")
|
||
return False
|
||
|
||
log("Connected")
|
||
session = Session(client)
|
||
await session.setup_refresh_subscription()
|
||
|
||
last_poll = 0.0 # D-03: poll immediately on first connect
|
||
used_successfully = False
|
||
consecutive_failures = 0 # D-03: zombie-link break counter
|
||
try:
|
||
while client.is_connected and not stop_event.is_set():
|
||
now = time.time()
|
||
elapsed = now - last_poll
|
||
if session.refresh_requested.is_set() or elapsed >= POLL_INTERVAL:
|
||
session.refresh_requested.clear()
|
||
|
||
# Local token usage (Session screen) needs no network — compute it
|
||
# every cycle so the watch keeps updating even with no/expired token.
|
||
payload = compute_today_usage()
|
||
|
||
# Rate-limit utilization is best-effort. A genuine 401/403 flags the
|
||
# token; a transient failure (network/DNS/5xx) leaves the tray state
|
||
# alone (SC#5: a DNS blip must not read as "token expired").
|
||
rl = None
|
||
auth_problem = False
|
||
|
||
# Self-refresh the OAuth token before using it. In a managed
|
||
# (Agent SDK) environment `claude login` is unavailable, so the
|
||
# daemon renews its own access token from the refresh token —
|
||
# otherwise the rate-limit % goes dark forever once it expires.
|
||
try:
|
||
await refresh_token_if_needed()
|
||
except Exception as e: # belt-and-braces: never crash the loop
|
||
log(f"Token refresh skipped: {e!r}")
|
||
|
||
token = read_token() # D-09: fresh each cycle (post-refresh)
|
||
if not token:
|
||
auth_problem = True
|
||
log("No token; sending local usage only")
|
||
else:
|
||
try:
|
||
rl = await poll_api(token)
|
||
except AuthError:
|
||
# Rejected despite the proactive refresh (e.g. expiresAt
|
||
# was stale/missing so we skipped it). Force one refresh
|
||
# and retry the poll once before flagging the token bad.
|
||
try:
|
||
forced = await refresh_token_if_needed(force=True)
|
||
except Exception as e:
|
||
log(f"Forced token refresh failed: {e!r}")
|
||
forced = False
|
||
if forced and (token := read_token()):
|
||
try:
|
||
rl = await poll_api(token)
|
||
except AuthError:
|
||
auth_problem = True
|
||
else:
|
||
auth_problem = True
|
||
|
||
if rl is not None:
|
||
payload.update(rl)
|
||
payload["ok"] = True # rate-limit data is fresh
|
||
else:
|
||
payload["ok"] = False # rate-limit unknown -> watch usage view goes idle
|
||
|
||
if await session.write_payload(payload):
|
||
last_poll = time.time()
|
||
used_successfully = True
|
||
consecutive_failures = 0 # D-03: reset on success
|
||
if rl is not None:
|
||
if tray_state:
|
||
tray_state.set_connected(time.time())
|
||
elif auth_problem:
|
||
if tray_state:
|
||
tray_state.set_error("token expired — run claude login")
|
||
# transient rate-limit failure: leave tray state unchanged
|
||
else:
|
||
consecutive_failures += 1
|
||
if consecutive_failures >= ZOMBIE_BREAK_LIMIT:
|
||
log(
|
||
f"Zombie link detected ({consecutive_failures} consecutive"
|
||
f" write failures); abandoning connection"
|
||
)
|
||
break
|
||
|
||
# Wake on a refresh request OR a stop, whichever comes first. Waking
|
||
# promptly on stop_event is what lets the finally below run
|
||
# client.disconnect() before the process exits, so the peer gets a
|
||
# clean GATT disconnect (returns to its waiting screen) instead of
|
||
# being left frozen on stale data after Quit (SC#3 graceful shutdown).
|
||
await _wait_first(session.refresh_requested, stop_event, timeout=TICK)
|
||
finally:
|
||
# Clean GATT disconnect on the way out — this is what tells the peripheral
|
||
# the link is gone. WinRT can surface a raw OSError (not BleakError) here,
|
||
# so swallow both; the link tears down regardless once we exit.
|
||
try:
|
||
await client.disconnect()
|
||
except (BleakError, OSError):
|
||
pass
|
||
|
||
log("Device disconnected" if not stop_event.is_set() else "Stopping")
|
||
return used_successfully
|
||
|
||
|
||
def _next_backoff(current: int, cap: int) -> int:
|
||
"""D-05: double current backoff value, clamped to cap.
|
||
|
||
Pure helper — unit-testable without driving the main loop.
|
||
Used by both slow-search (cap=60) and fast-reconnect (cap=RECONNECT_BACKOFF_CAP) regimes.
|
||
"""
|
||
return min(current * 2, cap)
|
||
|
||
|
||
async def main(tray_state=None) -> None:
|
||
stop_event = asyncio.Event()
|
||
loop = asyncio.get_running_loop()
|
||
|
||
# Populate the shared state object so the tray can route Quit through
|
||
# loop.call_soon_threadsafe (RESEARCH Pitfall 2). Additive — the existing
|
||
# stop_event = asyncio.Event() line above is unchanged.
|
||
if tray_state is not None:
|
||
tray_state.loop = loop
|
||
tray_state.stop_event = stop_event
|
||
|
||
def _stop(*_args: object) -> None:
|
||
log("Daemon stopping")
|
||
stop_event.set()
|
||
|
||
# OS signal handlers can only be installed from the main thread, and
|
||
# loop.add_signal_handler is unsupported on Windows. When running under the
|
||
# tray (04-03) the loop lives in a background thread and the tray owns clean
|
||
# shutdown via stop_event (loop.call_soon_threadsafe), so skip silently there.
|
||
if threading.current_thread() is threading.main_thread():
|
||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||
try:
|
||
loop.add_signal_handler(sig, _stop)
|
||
except NotImplementedError:
|
||
# Windows: add_signal_handler not supported; fall back to signal.signal
|
||
try:
|
||
signal.signal(sig, _stop)
|
||
except ValueError:
|
||
# Not the main thread of the main interpreter — tray owns shutdown.
|
||
pass
|
||
|
||
log("=== Claude Usage Tracker Daemon (BLE, Windows) ===")
|
||
log(f"Poll interval: {POLL_INTERVAL}s")
|
||
|
||
# D-05: two distinct backoff regimes — slow-search (device absent) vs fast-reconnect (link dropped)
|
||
search_backoff = 1 # caps at 60s — gentle, for a device that is genuinely absent/off
|
||
reconnect_backoff = 1 # caps at RECONNECT_BACKOFF_CAP — fast, to clear the 120s SLA after a drop
|
||
while not stop_event.is_set():
|
||
try:
|
||
device = await scan_for_device()
|
||
if not device:
|
||
# Slow-search regime: device was not found by scan — back off gently
|
||
if tray_state:
|
||
tray_state.set_scanning()
|
||
log(f"Device not found, retrying in {search_backoff}s...")
|
||
try:
|
||
await asyncio.wait_for(stop_event.wait(), timeout=search_backoff)
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
search_backoff = _next_backoff(search_backoff, 60)
|
||
continue
|
||
|
||
ok = await connect_and_run(device, stop_event, tray_state)
|
||
if not ok:
|
||
# Fast-reconnect regime: had/attempted a link that dropped — retry quickly
|
||
if tray_state:
|
||
tray_state.set_scanning()
|
||
log(f"Connection lost, reconnecting in {reconnect_backoff}s...")
|
||
try:
|
||
await asyncio.wait_for(stop_event.wait(), timeout=reconnect_backoff)
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
reconnect_backoff = _next_backoff(reconnect_backoff, RECONNECT_BACKOFF_CAP)
|
||
else:
|
||
# Successful session — reset reconnect counter to floor; search_backoff also reset
|
||
reconnect_backoff = 1
|
||
search_backoff = 1
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception as e:
|
||
# The whole BLE adapter can disappear (USB dongle unplugged) — bleak/
|
||
# WinRT then raises from the scan or the connect. Never let that kill
|
||
# the loop: log, show Scanning, back off, and keep retrying so
|
||
# replugging the adapter recovers on its own, no manual restart
|
||
# (field SC: pulled the dongle -> daemon crashed and stayed down).
|
||
if tray_state:
|
||
tray_state.set_scanning()
|
||
log(f"BLE error ({type(e).__name__}: {e}); retrying in {search_backoff}s...")
|
||
try:
|
||
await asyncio.wait_for(stop_event.wait(), timeout=search_backoff)
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
search_backoff = _next_backoff(search_backoff, 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
if sys.platform != "win32":
|
||
print(
|
||
"Warning: running under Linux/WSL — WinRT BLE will not be available.",
|
||
file=sys.stderr,
|
||
)
|
||
try:
|
||
asyncio.run(main())
|
||
except KeyboardInterrupt:
|
||
sys.exit(0)
|