From 63bcdcaa43dc09eba605ea15517aa21a1dd6887f Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Fri, 18 Sep 2026 23:34:56 -0700 Subject: [PATCH] fix(auth): refresh an expired Codex token before the quota-restored probe _probe_codex_quota_restored was called with the exhausted pool entry's stored access token. Exhausted entries are skipped by the proactive refresh chain, so for any cooldown longer than the access-token lifetime (the weekly-cap case the probe exists for) that token has expired: the usage endpoint answers 401 token_expired, the probe maps it to None, and the cooldown is kept. A mid-cooldown credit top-up or plan upgrade was therefore never discovered until last_error_reset_at elapsed, even though the refresh token was still valid (#89415). Add _refresh_expired_codex_probe_token: when the stored token is expired and a refresh token exists, rotate the pair via refresh_codex_oauth_pure (cooldown untouched), persist it (refresh tokens are single-use), and probe with the live token. Both probe callers use it: the pool-only resolve path (_probe_codex_pool_entry_quota_restored) and CredentialPool._codex_quota_restored_upstream. A probe that still reports 100% keeps the bench; the rotated tokens are persisted either way. Fixes #89415 --- agent/credential_pool.py | 8 ++ hermes_cli/auth.py | 5 +- hermes_cli/auth_codex.py | 48 ++++++++++- .../hermes_cli/test_auth_codex_quota_probe.py | 83 +++++++++++++++++++ 4 files changed, 139 insertions(+), 5 deletions(-) diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 6ea643dabb..bff676c97f 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -1670,6 +1670,14 @@ class CredentialPool(CredentialPoolAdminMixin, CredentialPoolModelCooldownMixin) if not token: return False try: + # An exhausted entry is skipped by the refresh chain, so its stored token is usually + # expired by probe time (401 -> None -> cooldown kept, #89415): refresh it first. + fresh = auth_mod._refresh_expired_codex_probe_token(token, entry.refresh_token) + if fresh: + entry = self._adopt( + entry, access_token=fresh["access_token"], refresh_token=fresh["refresh_token"], + ) + token = entry.access_token or token return bool(auth_mod._probe_codex_quota_restored(token, base_url=entry.base_url)) except Exception: logger.debug("Codex quota-restored probe failed", exc_info=True) diff --git a/hermes_cli/auth.py b/hermes_cli/auth.py index 28b31563a9..5b121a0f61 100644 --- a/hermes_cli/auth.py +++ b/hermes_cli/auth.py @@ -76,8 +76,9 @@ from hermes_cli.auth_codex import ( # noqa: F401 re-exported _codex_access_token_is_expiring, _codex_device_code_login, _codex_http_client, _codex_pool_rate_limit_status, _codex_quota_probe_cache, _codex_usage_probe_url, _import_codex_cli_tokens, _is_codex_rate_limit_shaped, _login_openai_codex, - _probe_codex_quota_restored, _read_codex_tokens, _refresh_codex_auth_tokens, _save_codex_tokens, - clear_codex_pool_quota_cooldowns, refresh_codex_oauth_pure, resolve_codex_runtime_credentials) + _probe_codex_quota_restored, _read_codex_tokens, _refresh_codex_auth_tokens, + _refresh_expired_codex_probe_token, _save_codex_tokens, clear_codex_pool_quota_cooldowns, + refresh_codex_oauth_pure, resolve_codex_runtime_credentials) from hermes_cli.auth_spotify import ( # noqa: F401 re-exported _refresh_spotify_oauth_state, get_spotify_auth_status, login_spotify_command, resolve_spotify_runtime_credentials) diff --git a/hermes_cli/auth_codex.py b/hermes_cli/auth_codex.py index db70044d89..68ed44ff99 100644 --- a/hermes_cli/auth_codex.py +++ b/hermes_cli/auth_codex.py @@ -557,9 +557,7 @@ def resolve_codex_runtime_credentials( # Before surfacing the persisted cooldown, ask the usage endpoint whether the quota # reset early (banked reset redeemed, plan upgraded): ``last_error_reset_at`` can be # days in the future while the account is already usable again. - stale_token = _stripped(pool_rate_limit.get("access_token")) - if stale_token and _probe_codex_quota_restored( - stale_token, base_url=pool_rate_limit.get("base_url")): + if _probe_codex_pool_entry_quota_restored(pool_rate_limit): logger.info("Codex quota restored upstream — clearing stale pool cooldown(s).") clear_codex_pool_quota_cooldowns() pool_token = _pool_codex_access_token() @@ -683,6 +681,49 @@ def _probe_codex_quota_restored( return result +def _refresh_expired_codex_probe_token( + access_token: Any, refresh_token: Any) -> Optional[Dict[str, Any]]: + """Refresh an EXPIRED stored access token so the quota probe can get a real answer. + + Exhausted pool entries are skipped by the proactive refresh chain (#44799), so by the time + anything probes with the stored token it has expired; the usage endpoint answers + ``401 token_expired``, the probe returns None, and the cooldown is kept until + ``last_error_reset_at`` no matter what happened upstream (top-up, plan upgrade) — #89415. + Returns the rotated token pair (callers MUST persist it: refresh tokens are single-use) or + None when no refresh was needed/possible. The cooldown itself is left untouched. + """ + token, refresh = _stripped(access_token), _stripped(refresh_token) + if not token or not refresh or not _codex_access_token_is_expiring(token, 0): + return None + try: + return refresh_codex_oauth_pure(token, refresh) + except Exception: + logger.debug("Codex pre-probe token refresh failed", exc_info=True) + return None + + +def _probe_codex_pool_entry_quota_restored(entry: Dict[str, Any]) -> Optional[bool]: + """``_probe_codex_quota_restored`` for a persisted pool entry, refreshing an expired token first.""" + from hermes_cli.auth import _auth_store_lock, _load_auth_store, _save_auth_store + token = _stripped(entry.get("access_token")) + fresh = _refresh_expired_codex_probe_token(token, entry.get("refresh_token")) + if fresh: + token = fresh["access_token"] + try: + with _auth_store_lock(): + auth_store = _load_auth_store() + for disk_entry in _codex_pool_dicts(_pool_entries(auth_store, "openai-codex")): + if disk_entry.get("id") == entry.get("id"): + disk_entry.update(fresh) + _save_auth_store(auth_store) + break + except Exception: + logger.debug("Failed to persist refreshed Codex pool tokens", exc_info=True) + if not token: + return None + return _probe_codex_quota_restored(token, base_url=entry.get("base_url")) + + def clear_codex_pool_quota_cooldowns(access_token: Optional[str] = None) -> int: """Clear rate-limit cooldowns on persisted openai-codex pool entries. @@ -732,6 +773,7 @@ def _codex_pool_rate_limit_status() -> Optional[Dict[str, Any]]: "label": entry.get("label"), "last_refresh": entry.get("last_refresh"), "reset_at": reset_at, "reason": entry.get("last_error_reason"), "message": entry.get("last_error_message"), "access_token": token.strip(), + "refresh_token": entry.get("refresh_token"), "id": entry.get("id"), "base_url": entry.get("base_url")} except Exception: logger.debug("Codex pool rate-limit lookup failed", exc_info=True) diff --git a/tests/hermes_cli/test_auth_codex_quota_probe.py b/tests/hermes_cli/test_auth_codex_quota_probe.py index d319227480..115cafc12d 100644 --- a/tests/hermes_cli/test_auth_codex_quota_probe.py +++ b/tests/hermes_cli/test_auth_codex_quota_probe.py @@ -316,3 +316,86 @@ def test_pool_probe_not_fired_for_non_quota_exhaustion(tmp_path, monkeypatch): # --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# #89415 — the mid-cooldown probe must refresh an expired stored token first +# --------------------------------------------------------------------------- + + +def _expired_jwt_pool_store(now): + store = _pool_only_rate_limited_store(now) + entry = store["credential_pool"]["openai-codex"][0] + entry["access_token"] = _jwt({"exp": now - 7200}) # expired hours ago + entry["refresh_token"] = "rf-old" + return store + + +class _ExpiryAwareClient(_StubClient): + """Behaves like the real usage endpoint: an expired bearer gets 401 token_expired.""" + + def get(self, url, headers=None): + token = (headers or {}).get("Authorization", "").removeprefix("Bearer ") + if auth_codex._codex_access_token_is_expiring(token, 0): + self._calls.append({"url": url, "headers": dict(headers or {})}) + return _StubResponse(401, {"error": {"code": "token_expired"}}) + return super().get(url, headers=headers) + + +def _patch_expiry_aware_httpx(monkeypatch, response): + calls: list = [] + monkeypatch.setattr( + auth_mod.httpx, "Client", lambda **kwargs: _ExpiryAwareClient(calls, response) + ) + return calls + + +def _fake_refresh(monkeypatch, fresh_token, calls): + def _refresh(access_token, refresh_token, **kw): + calls.append(refresh_token) + return {"access_token": fresh_token, "refresh_token": "rf-new", "last_refresh": "now"} + + monkeypatch.setattr(auth_codex, "refresh_codex_oauth_pure", _refresh) + + +def test_resolver_refreshes_expired_token_before_probe(tmp_path, monkeypatch): + """Exhausted entries are skipped by the refresh chain, so the stored access token has + expired by the time the probe runs: /usage answers 401 -> None -> cooldown kept forever, + even after a top-up / plan upgrade. Refresh (keeping the cooldown) and probe live.""" + now = time.time() + hermes_home = tmp_path / "hermes" + _write_auth_store(hermes_home, _expired_jwt_pool_store(now)) + monkeypatch.setenv("HERMES_HOME", str(hermes_home)) + fresh = _jwt({"exp": now + 3600}) + refresh_calls: list = [] + _fake_refresh(monkeypatch, fresh, refresh_calls) + http_calls = _patch_expiry_aware_httpx(monkeypatch, _StubResponse(200, _usage_payload(0.0, 0.0))) + + resolved = resolve_codex_runtime_credentials() + + assert refresh_calls == ["rf-old"] + assert http_calls[0]["headers"]["Authorization"] == f"Bearer {fresh}" + assert resolved["api_key"] == fresh + entry = json.loads((hermes_home / "auth.json").read_text())["credential_pool"]["openai-codex"][0] + assert entry["refresh_token"] == "rf-new" + assert entry["last_status"] is None + + +def test_resolver_keeps_cooldown_when_fresh_probe_still_exhausted(tmp_path, monkeypatch): + """Control: refresh succeeds, live probe still says 100% -> cooldown stays, but the + rotated (single-use) token pair is persisted so the grant is not burned.""" + now = time.time() + hermes_home = tmp_path / "hermes" + _write_auth_store(hermes_home, _expired_jwt_pool_store(now)) + monkeypatch.setenv("HERMES_HOME", str(hermes_home)) + fresh = _jwt({"exp": now + 3600}) + _fake_refresh(monkeypatch, fresh, []) + _patch_expiry_aware_httpx(monkeypatch, _StubResponse(200, _usage_payload(0.0, 100.0))) + + with pytest.raises(AuthError, match="quota exhausted"): + resolve_codex_runtime_credentials() + + entry = json.loads((hermes_home / "auth.json").read_text())["credential_pool"]["openai-codex"][0] + assert entry["access_token"] == fresh + assert entry["refresh_token"] == "rf-new" + assert entry["last_status"] == "exhausted"