fix(codex): run the refresh inside the source store's transaction so shared-root peers adopt, not replay

Follow-up to #110024 (ehz0ah's review thread). `_refresh_codex_auth_tokens` POSTed the
single-use refresh token to OpenAI first and only then entered
`_provider_state_transaction("openai-codex")` for the write-back, so root's lock covered the
save alone. `resolve_codex_runtime_credentials` holds only the caller's own profile lock, so
two profiles borrowing the same ROOT grant could both submit `old-rt`; last root save won and
OpenAI answered `refresh_token_reused` / revoked the family — the failure #87503 exists to
prevent.

The transaction now spans re-read -> endpoint refresh -> write-back:
- enter `_provider_state_transaction` first; the yielded state is root's, re-read under root's
  lock. If its refresh token already differs from the one we were about to submit, a peer
  rotated it: adopt the stored pair and return without touching the endpoint.
- otherwise POST and write back through `_store_codex_tokens_in`, the body of
  `_save_codex_tokens` split out so it can run inside an already-open transaction.
  `_save_codex_tokens` keeps its signature for the login/import/CLI-recovery callers.

Holding the advisory flock across the network call is safe here and already the established
shape: `resolve_codex_runtime_credentials` holds the active-store lock across the same POST,
and every waiter's timeout is `max(AUTH_LOCK_TIMEOUT_SECONDS, refresh_timeout + 5)`, i.e. it
outlives one full endpoint timeout. `_load_auth_store` readers never take the lock, so
readers are not blocked; `_file_lock` is reentrant per thread per path, so the nested
transaction inside the caller's lock and the CLI-recovery save inside the transaction both
re-enter cleanly. A release-POST-retake variant would reopen the window it is meant to close.

Test: two refreshers with the same stale pre-read pair against a rotate-once endpoint that
rejects any replay — the endpoint sees `old-rt` exactly once, both callers end with the
rotated pair, root holds it, the profile store stays unshadowed. Red on origin/main
(`refresh_token_reused` surfaces for the second caller).
This commit is contained in:
kshitijk4poor
2026-09-14 20:11:56 +05:30
committed by kshitij
parent a19160bbf0
commit e117e792b6
2 changed files with 132 additions and 44 deletions

View File

@@ -157,29 +157,40 @@ def _save_codex_tokens(
Only a token REFRESH passes ``write_through=True``: a fresh login or import under a profile
is the profile's own grant and must not overwrite the root account it was borrowing.
"""
from hermes_cli.auth import _provider_state_transaction
with _provider_state_transaction("openai-codex") as (auth_store, state, source_path):
_store_codex_tokens_in(
auth_store, state, source_path, tokens, last_refresh, label, write_through=write_through)
def _store_codex_tokens_in(
auth_store: Dict[str, Any], state: Optional[Dict[str, Any]], source_path: Optional[Path],
tokens: Dict[str, str], last_refresh: Optional[str], label: Optional[str] = None, *,
write_through: bool,
) -> None:
"""Body of ``_save_codex_tokens`` for a caller already inside ``_provider_state_transaction``."""
from hermes_cli.auth import (
_auth_file_path, _load_auth_store, _provider_state_transaction, _same_path,
_save_auth_store, _store_provider_state, _utc_now_z)
_auth_file_path, _load_auth_store, _same_path, _save_auth_store, _store_provider_state,
_utc_now_z)
if last_refresh is None:
last_refresh = _utc_now_z()
with _provider_state_transaction("openai-codex") as (auth_store, state, source_path):
state = dict(state) if state else {}
# Capture the previous singleton tokens BEFORE overwriting: the pool sync uses them to
# tell legacy singleton-aliases (refresh) from independent ``auth add`` accounts (keep).
previous_singleton_tokens = (
state.get("tokens") if isinstance(state.get("tokens"), dict) else None)
state.update(tokens=tokens, last_refresh=last_refresh, auth_mode="chatgpt")
if label and str(label).strip():
state["label"] = str(label).strip()
target_store, target_path, set_active = auth_store, None, True
if write_through and source_path is not None and not _same_path(source_path, _auth_file_path()):
# Root-borrowed grant: the transaction already holds root's lock, so write the rotated
# chain into ROOT's store (never set_active — a refresh is not a provider choice).
target_store, target_path, set_active = _load_auth_store(source_path), source_path, False
_store_provider_state(target_store, "openai-codex", state, set_active=set_active)
_sync_codex_pool_entries(
target_store, tokens, last_refresh, previous_singleton_tokens=previous_singleton_tokens)
_save_auth_store(target_store, target_path=target_path)
state = dict(state) if state else {}
# Capture the previous singleton tokens BEFORE overwriting: the pool sync uses them to
# tell legacy singleton-aliases (refresh) from independent ``auth add`` accounts (keep).
previous_singleton_tokens = (
state.get("tokens") if isinstance(state.get("tokens"), dict) else None)
state.update(tokens=tokens, last_refresh=last_refresh, auth_mode="chatgpt")
if label and str(label).strip():
state["label"] = str(label).strip()
target_store, target_path, set_active = auth_store, None, True
if write_through and source_path is not None and not _same_path(source_path, _auth_file_path()):
# Root-borrowed grant: the transaction already holds root's lock, so write the rotated
# chain into ROOT's store (never set_active — a refresh is not a provider choice).
target_store, target_path, set_active = _load_auth_store(source_path), source_path, False
_store_provider_state(target_store, "openai-codex", state, set_active=set_active)
_sync_codex_pool_entries(
target_store, tokens, last_refresh, previous_singleton_tokens=previous_singleton_tokens)
_save_auth_store(target_store, target_path=target_path)
def _recover_codex_tokens_from_cli(reason: str) -> Optional[Dict[str, str]]:
@@ -374,29 +385,46 @@ def refresh_codex_oauth_pure(
def _refresh_codex_auth_tokens(tokens: Dict[str, str], timeout_seconds: float) -> Dict[str, str]:
"""Refresh Codex access token using the refresh token."""
from hermes_cli.auth import _save_codex_tokens, refresh_codex_oauth_pure
try:
refreshed = refresh_codex_oauth_pure(
str(tokens.get("access_token", "") or ""), str(tokens.get("refresh_token", "") or ""),
timeout_seconds=timeout_seconds)
except AuthError as exc:
# Self-heal cross-store rotation: refresh_tokens are single-use, so when the Codex CLI (or
# another Hermes process) rotates the shared token this frozen copy fails with a
# relogin-required error (invalid_grant / refresh_token_reused / 401). Adopt the canonical
# fresh token from ~/.codex/auth.json before surfacing a hard 401. Transient failures
# (429 quota) keep relogin_required=False — the stored token is still valid — re-raise.
if not getattr(exc, "relogin_required", False):
raise
imported = _recover_codex_tokens_from_cli(
f"refresh_token rejected: {getattr(exc, 'code', None) or 'auth_error'}")
if not imported:
raise
return imported
updated_tokens = {
**tokens, "access_token": refreshed["access_token"],
"refresh_token": refreshed["refresh_token"]}
_save_codex_tokens(updated_tokens, write_through=True)
"""Refresh Codex access token using the refresh token.
The whole re-read -> endpoint POST -> write-back runs inside the SOURCE store's transaction:
two profiles borrowing the same root grant otherwise both submit the same single-use refresh
token (each holds only its own profile lock) and OpenAI revokes the family. A waiter that
finds root already rotated by its peer adopts the stored pair instead of replaying the
consumed token. The caller's lock timeout already covers a full endpoint timeout.
"""
from hermes_cli.auth import _provider_state_transaction, refresh_codex_oauth_pure
with _provider_state_transaction("openai-codex") as (auth_store, state, source_path):
stored = (state or {}).get("tokens")
stored = stored if isinstance(stored, dict) else {}
stored_rt = _stripped(stored.get("refresh_token"))
if stored_rt and stored_rt != _stripped(tokens.get("refresh_token")):
logger.info("Codex refresh token already rotated by a peer — adopting the stored pair.")
return {**tokens, "access_token": _stripped(stored.get("access_token")),
"refresh_token": stored_rt}
try:
refreshed = refresh_codex_oauth_pure(
str(tokens.get("access_token", "") or ""), str(tokens.get("refresh_token", "") or ""),
timeout_seconds=timeout_seconds)
except AuthError as exc:
# Self-heal cross-store rotation: refresh_tokens are single-use, so when the Codex CLI
# (or another Hermes process) rotates the shared token this frozen copy fails with a
# relogin-required error (invalid_grant / refresh_token_reused / 401). Adopt the
# canonical fresh token from ~/.codex/auth.json before surfacing a hard 401. Transient
# failures (429 quota) keep relogin_required=False — the stored token is still valid —
# re-raise.
if not getattr(exc, "relogin_required", False):
raise
imported = _recover_codex_tokens_from_cli(
f"refresh_token rejected: {getattr(exc, 'code', None) or 'auth_error'}")
if not imported:
raise
return imported
updated_tokens = {
**tokens, "access_token": refreshed["access_token"],
"refresh_token": refreshed["refresh_token"]}
_store_codex_tokens_in(
auth_store, state, source_path, updated_tokens, None, write_through=True)
return updated_tokens

View File

@@ -7,11 +7,13 @@ Token values are synthetic placeholders.
"""
import json
import threading
from pathlib import Path
import httpx
import pytest
from hermes_cli import auth
from hermes_cli import auth, auth_codex
def _pair(prefix: str) -> dict:
@@ -81,3 +83,61 @@ def test_profile_owned_grant_stays_local(profile_env):
auth._save_codex_tokens(_pair("login"))
assert _read(profile_path)["providers"]["openai-codex"]["tokens"] == _pair("login")
assert _read(root_path)["providers"]["openai-codex"]["tokens"] == _pair("root")
class _RotatingEndpoint:
"""Token endpoint that rotates ``old-rt`` once and rejects any replay of a consumed token."""
def __init__(self, hold_seconds: float):
self.hold_seconds = hold_seconds
self.seen: list = []
self._guard = threading.Lock()
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def post(self, url, *, headers=None, data=None):
import time
with self._guard:
replay = data["refresh_token"] in self.seen
self.seen.append(data["refresh_token"])
time.sleep(self.hold_seconds) # a concurrent refresher must wait on the lock, not overlap
if replay:
return httpx.Response(400, json={"error": "refresh_token_reused"})
return httpx.Response(200, json={"access_token": "new-at", "refresh_token": "new-rt"})
def test_concurrent_refreshes_of_shared_root_grant_submit_old_token_once(profile_env, monkeypatch):
"""Two refreshers holding the same stale pre-read pair: the second re-reads ROOT under its lock,
sees the peer's rotated pair and adopts it instead of replaying the single-use token."""
profile_path, root_path = profile_env
_write(root_path, {
"version": 1,
"providers": {"openai-codex": {"auth_mode": "chatgpt", "tokens": _pair("old")}},
})
_write(profile_path, {"version": 1, "providers": {}})
endpoint = _RotatingEndpoint(hold_seconds=0.3)
monkeypatch.setattr(auth_codex, "_codex_http_client", lambda **kw: endpoint)
results, errors = {}, {}
def _refresh(name):
try:
results[name] = auth._refresh_codex_auth_tokens(_pair("old"), timeout_seconds=5.0)
except Exception as exc: # pragma: no cover - surfaced via the assertion below
errors[name] = exc
workers = [threading.Thread(target=_refresh, args=(n,)) for n in ("a", "b")]
for w in workers:
w.start()
for w in workers:
w.join(timeout=10)
assert errors == {}
assert endpoint.seen == ["old-rt"]
assert results == {"a": _pair("new"), "b": _pair("new")}
assert _read(root_path)["providers"]["openai-codex"]["tokens"] == _pair("new")
assert "openai-codex" not in _read(profile_path).get("providers", {})