diff --git a/daemon/claude_usage_daemon_windows.py b/daemon/claude_usage_daemon_windows.py index 3ab9fec..905023a 100644 --- a/daemon/claude_usage_daemon_windows.py +++ b/daemon/claude_usage_daemon_windows.py @@ -57,6 +57,17 @@ API_BODY = { "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. @@ -416,6 +427,106 @@ def _read_expiry() -> str: 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. @@ -504,7 +615,17 @@ async def connect_and_run(device, stop_event: asyncio.Event, tray_state=None) -> # alone (SC#5: a DNS blip must not read as "token expired"). rl = None auth_problem = False - token = read_token() # D-09: fresh each cycle + + # 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") @@ -512,7 +633,21 @@ async def connect_and_run(device, stop_event: asyncio.Event, tray_state=None) -> try: rl = await poll_api(token) except AuthError: - auth_problem = True + # 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)