hermes-admin/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 33 Packages Projects Releases Wiki Activity

fix(auth): refresh an expired Codex token before the quota-restored probe

Browse Source 
_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
This commit is contained in:
Teknium
2026-09-18 23:34:56 -07:00
parent 8530d81b55
commit 63bcdcaa43
4 changed files with 139 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
agent/credential_pool.py
View File

@@ -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)

5
hermes_cli/auth.py
View File

@@ -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)

48
hermes_cli/auth_codex.py
View File

@@ -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)

83
tests/hermes_cli/test_auth_codex_quota_probe.py
View File

@@ -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"