fix(credential_pool): use source-path-based write-through to root (#74339)

_sync_device_code_entry_to_auth_store used key-presence on the profile
store to decide whether to write-through rotated tokens to the global
root.  _store_provider_state unconditionally creates that key, so every
refresh after the first self-disabled the write-through — root kept a
revoked refresh token and every other profile died with
refresh_token_reused / invalid_grant.

Fix: use _load_provider_state_with_source to learn where the grant was
resolved from.  When the source is the global root, write back only to
root and skip _store_provider_state so the profile never accrues a
shadowing providers.<id> key that blocks future root fallback.

Add regression test verifying write-through fires on refresh 2+, not
just the first call.
This commit is contained in:
praneshnikhar 2026-07-31 19:07:33 +05:30 committed by Teknium
parent 29eac371d1
commit 4e6299af48
2 changed files with 141 additions and 32 deletions

View File

@ -28,10 +28,13 @@ from hermes_cli.auth import (
_auth_store_lock,
_codex_access_token_is_expiring,
_decode_jwt_claims,
_global_auth_file_path,
_load_auth_store,
_load_provider_state,
_load_provider_state_with_source,
_resolve_kimi_base_url,
_resolve_zai_base_url,
_same_path,
_save_auth_store,
_save_provider_state,
_store_provider_state,
@ -1039,32 +1042,58 @@ class CredentialPool:
try:
with _auth_store_lock():
auth_store = _load_auth_store()
# Decide BEFORE writing whether this profile is reading the
# grant from the global root (no own providers.<id> block) vs.
# genuinely shadowing it. A pool refresh rotates single-use
# OAuth refresh tokens, so a profile that resolved the grant
# from root MUST write the rotated chain back to root too —
# otherwise root keeps a revoked refresh token and every other
# profile reading the stale root grant dies with
# refresh_token_reused / invalid_grant once its access token
# expires. This mirrors the xAI write-through in
# hermes_cli.auth._save_xai_oauth_tokens (#43589); the pool
# refresh path is the Codex/xAI analog reported in #48415.
_wt_provider_id = {
"nous": "nous",
"openai-codex": "openai-codex",
"xai-oauth": "xai-oauth",
}.get(self.provider)
write_through_to_root = bool(_wt_provider_id) and not (
isinstance(auth_store.get("providers"), dict)
and isinstance(
auth_store["providers"].get(_wt_provider_id), dict
)
)
# Resolve state and track which store it came from — the
# source path tells us whether this profile genuinely owns
# its provider block or is reading from the global root.
# #74339: the old key-presence check decided write-through
# on whether the profile had ``providers.<id>`` BEFORE the
# save — correct for the first refresh but self-sealing
# because ``_store_provider_state`` unconditionally creates
# that key inside the same function. Once the profile has
# the key, every subsequent refresh silently disables the
# root write-through and root keeps a revoked refresh token.
#
# Fix: use ``_load_provider_state_with_source`` to learn
# where the state was resolved from. When the grant was
# resolved from the global root, write back *only* to root
# and skip ``_store_provider_state`` for the profile so the
# profile does not accrue a shadowing ``providers.<id>``
# key that blocks both the root fallback and the write-through
# on subsequent calls.
if self.provider == "nous":
state = _load_provider_state(auth_store, "nous")
state, source_path = _load_provider_state_with_source(
auth_store, "nous"
)
if state is None:
return
elif self.provider == "openai-codex":
state, source_path = _load_provider_state_with_source(
auth_store, "openai-codex"
)
if not isinstance(state, dict):
return
elif self.provider == "xai-oauth":
state, source_path = _load_provider_state_with_source(
auth_store, "xai-oauth"
)
if not isinstance(state, dict):
return
else:
return
global_root = _global_auth_file_path()
is_from_root = bool(
source_path is not None
and global_root is not None
and _same_path(source_path, global_root)
)
if self.provider == "nous":
state["access_token"] = entry.access_token
if entry.refresh_token:
state["refresh_token"] = entry.refresh_token
@ -1082,12 +1111,8 @@ class CredentialPool:
state[extra_key] = val
if entry.inference_base_url:
state["inference_base_url"] = entry.inference_base_url
_store_provider_state(auth_store, "nous", state, set_active=False)
elif self.provider == "openai-codex":
state = _load_provider_state(auth_store, "openai-codex")
if not isinstance(state, dict):
return
tokens = state.get("tokens")
if not isinstance(tokens, dict):
return
@ -1096,12 +1121,8 @@ class CredentialPool:
tokens["refresh_token"] = entry.refresh_token
if entry.last_refresh:
state["last_refresh"] = entry.last_refresh
_store_provider_state(auth_store, "openai-codex", state, set_active=False)
elif self.provider == "xai-oauth":
state = _load_provider_state(auth_store, "xai-oauth")
if not isinstance(state, dict):
return
tokens = state.get("tokens")
if not isinstance(tokens, dict):
return
@ -1110,16 +1131,26 @@ class CredentialPool:
tokens["refresh_token"] = entry.refresh_token
if entry.last_refresh:
state["last_refresh"] = entry.last_refresh
_store_provider_state(auth_store, "xai-oauth", state, set_active=False)
else:
return
_save_auth_store(auth_store)
if write_through_to_root and _wt_provider_id:
if is_from_root and _wt_provider_id:
# Grant was resolved from root — write back to root
# only. Do NOT call _store_provider_state on the
# profile auth_store (it would create a shadowing
# providers.<id> key that disables write-through on
# the next refresh — #74339).
# _load_provider_state has root fallback, so the
# profile can always read fresh tokens from root
# without needing its own providers block.
_write_through_provider_state_to_global_root(
_wt_provider_id, state
)
else:
# Profile genuinely owns this provider — write to
# the profile store as normal.
_store_provider_state(
auth_store, self.provider, state, set_active=False
)
_save_auth_store(auth_store)
except Exception as exc:
logger.debug("Failed to sync %s pool entry back to auth store: %s", self.provider, exc)

View File

@ -233,3 +233,81 @@ def test_codex_pool_refresh_holds_auth_store_lock_across_post(monkeypatch, tmp_p
# The invariant: the single-use token POST ran inside the auth-store lock.
assert lock_held["during_post"] is True
def test_write_through_fires_on_every_refresh_not_just_first(
profile_and_root, monkeypatch
):
"""Write-through to root must fire on the 2nd, 3rd, … refresh too (#74339).
The old key-presence check decided write-through on whether the *profile*
store had ``providers.<id>`` BEFORE the save a key that
``_store_provider_state()`` unconditionally created. Net effect: first
refresh write-through fires; every later refresh silently disabled
because the profile now "owned" the block, even though it never
performed its own OAuth grant.
The fix skips ``_store_provider_state`` entirely when the grant was
resolved from root, so the profile never accrues a shadowing key and
``_load_provider_state_with_source`` always resolves from root.
"""
profile_path, root_path = profile_and_root
_write_store(
root_path,
{
"version": 1,
"providers": {
"openai-codex": {
"tokens": {"access_token": "root-ac", "refresh_token": "root-rf"}
}
},
},
)
provider = "openai-codex"
call_count = [0]
def counting_root_write(provider_id, state):
call_count[0] += 1
monkeypatch.setattr(
CP, "_write_through_provider_state_to_global_root", counting_root_write
)
# After patching A's module-level attributes, the bare-name imports in
# credential_pool.py still hold references to the original functions
# (``from X import Y`` creates a local binding that does not update when
# ``X.Y`` is reassigned). Patch CP's bindings separately so the
# ``_sync_device_code_entry_to_auth_store`` method — whose __globals__
# are ``agent.credential_pool.__dict__`` — sees the mocked paths.
monkeypatch.setattr(CP, "_global_auth_file_path", lambda: root_path)
monkeypatch.setattr(CP, "_same_path", lambda a, b: a == b)
# ---- REFRESH 1 ----
_write_store(profile_path, {"version": 1})
entry1 = _entry(
provider, id="c1", access_token="ac1", refresh_token="rf1"
)
pool1 = CredentialPool(provider, [entry1])
pool1._sync_device_code_entry_to_auth_store(entry1)
assert call_count[0] == 1, "refresh 1: write-through must fire (#74339)"
# After refresh 1 the profile should NOT have a providers.openai-codex
# block (the fix skipped _store_provider_state because the grant came
# from root). This prevents the self-sealing that broke refresh 2+.
profile_store = _read_store(profile_path)
assert "openai-codex" not in profile_store.get("providers", {}), (
"profile must NOT accrue a shadowing providers.<id> block when the "
"grant was resolved from root — that key would disable write-through "
"on the next refresh (#74339)"
)
# ---- REFRESH 2 (same scenario, rotated tokens) ----
entry2 = _entry(
provider, id="c2", access_token="ac2", refresh_token="rf2"
)
pool2 = CredentialPool(provider, [entry2])
pool2._sync_device_code_entry_to_auth_store(entry2)
assert call_count[0] == 2, (
"refresh 2: write-through must fire even after a prior sync-back. "
"The old code self-disabled here (#74339)"
)