From 4e6299af48a9b0a71bdad1a0dc4d750c770befbe Mon Sep 17 00:00:00 2001 From: praneshnikhar Date: Fri, 31 Jul 2026 19:07:33 +0530 Subject: [PATCH] fix(credential_pool): use source-path-based write-through to root (#74339) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _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. key that blocks future root fallback. Add regression test verifying write-through fires on refresh 2+, not just the first call. --- agent/credential_pool.py | 95 ++++++++++++------- ...test_credential_pool_oauth_writethrough.py | 78 +++++++++++++++ 2 files changed, 141 insertions(+), 32 deletions(-) diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 16857eba830dc..715f8b8b174c7 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -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. 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.`` 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.`` + # 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. 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) diff --git a/tests/agent/test_credential_pool_oauth_writethrough.py b/tests/agent/test_credential_pool_oauth_writethrough.py index 5794592802034..707b502d1c647 100644 --- a/tests/agent/test_credential_pool_oauth_writethrough.py +++ b/tests/agent/test_credential_pool_oauth_writethrough.py @@ -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.`` 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. 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)" + ) +