From bf7c7166485ec7fcba5ca3a63e9fd0c2d7c91e48 Mon Sep 17 00:00:00 2001 From: 686f6c61 Date: Wed, 5 Aug 2026 09:39:02 +0200 Subject: [PATCH] fix(agent): rebind pool entry id after env credential refresh Per-turn .env adoption could rewrite agent.api_key while leaving _credential_pool_entry_id on a previously rotated fallback. The next 429 then marked the healthy fallback exhausted via credential_id precedence (#79156). - Sync pool entry id after a successful env credential refresh - First look does not stomp a pool-rotated key with the env primary - mark_exhausted_and_rotate prefers api_key_hint when it disagrees with credential_id Fixes #79156 --- agent/credential_pool.py | 31 +++++++++++ run_agent.py | 24 ++++++++ tests/agent/test_credential_pool.py | 55 +++++++++++++++++++ .../test_env_credential_turn_refresh.py | 37 +++++++++++++ 4 files changed, 147 insertions(+) diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 85a6e79a75..072140134e 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -2045,6 +2045,37 @@ class CredentialPool: (e for e in self._entries if e.id == credential_id), None, ) + # #79156: when both identities are supplied and they disagree, + # trust the key that actually made the request. A stale + # ``_credential_pool_entry_id`` (e.g. after per-turn env + # refresh rewrote ``api_key`` without rebinding the id) would + # otherwise quarantine a healthy fallback for days. + if ( + entry is not None + and api_key_hint + and entry.runtime_api_key != api_key_hint + ): + hint_entry = next( + ( + e + for e in self._entries + if e.runtime_api_key == api_key_hint + ), + None, + ) + if hint_entry is not None: + logger.info( + "credential pool: credential_id %s runtime key " + "does not match api_key_hint; attributing failure " + "to key-matched entry %s instead (#79156)", + (entry.label or entry.id[:8]), + (hint_entry.label or hint_entry.id[:8]), + ) + entry = hint_entry + else: + # Id is stale and the request key is not in the pool — + # drop the id so we do not mark the wrong entry. + entry = None if entry is None and api_key_hint: # Prefer the specific entry whose API key matches the one that # actually failed. When this pool was freshly loaded from disk diff --git a/run_agent.py b/run_agent.py index 0aa009fae4..47b4a0e80e 100644 --- a/run_agent.py +++ b/run_agent.py @@ -5611,6 +5611,17 @@ class AIAgent: adopt = current_base == default_base and not ( base_url == current_base and api_key == self.api_key ) + # #79156: if the session already holds a pool-rotated key, do + # not treat that divergence as a boot-time env adoption. First + # look would otherwise stomp the rotated key with the env value + # while leaving ``_credential_pool_entry_id`` on the fallback. + if ( + adopt + and api_key != self.api_key + and getattr(self, "_credential_pool", None) is not None + and getattr(self, "_credential_pool_entry_id", None) + ): + adopt = False else: # Env unchanged → no-op; any drift from self.* is rotation/ # failover or config precedence — leave it alone. An edit is @@ -5654,6 +5665,19 @@ class AIAgent: self._client_kwargs.update(prior_client_kwargs) return False + # Rebind the pool entry id to the key we just adopted. Leaving a + # stale id after a key rewrite makes mark_exhausted_and_rotate + # quarantine the wrong credential on the next 429 (#79156). + try: + from agent.agent_runtime_helpers import sync_credential_pool_entry_id + + sync_credential_pool_entry_id(self) + except Exception: + logger.debug( + "sync_credential_pool_entry_id after env refresh failed", + exc_info=True, + ) + self._env_creds_seen = resolved logger.info( "Applied updated .env credentials for %s: endpoint %s", diff --git a/tests/agent/test_credential_pool.py b/tests/agent/test_credential_pool.py index 8d7e78c74e..72d2f04d01 100644 --- a/tests/agent/test_credential_pool.py +++ b/tests/agent/test_credential_pool.py @@ -143,6 +143,61 @@ def test_billing_rotation_marks_all_entries_sharing_failed_key(tmp_path, monkeyp assert statuses["cred-model-config"] == STATUS_EXHAUSTED +def test_stale_credential_id_prefers_api_key_hint(tmp_path, monkeypatch): + """#79156: disagreeing credential_id + api_key_hint must mark the key. + + After per-turn env refresh rewrites ``api_key`` without rebinding the + pool entry id, recovery still passes the stale id of the healthy + fallback together with the primary key that actually failed. The + healthy key must not inherit the primary's 429. + """ + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + monkeypatch.setattr("agent.anthropic_adapter.read_claude_code_credentials", lambda: None) + _write_auth_store( + tmp_path, + { + "version": 1, + "credential_pool": { + "anthropic": [ + { + "id": "cred-primary", + "label": "primary", + "auth_type": "api_key", + "priority": 0, + "source": "manual", + "access_token": "sk-ant-api-primary", + }, + { + "id": "cred-backup", + "label": "backup", + "auth_type": "api_key", + "priority": 1, + "source": "manual", + "access_token": "sk-ant-api-backup", + }, + ] + }, + }, + ) + + from agent.credential_pool import load_pool, STATUS_EXHAUSTED + + pool = load_pool("anthropic") + next_entry = pool.mark_exhausted_and_rotate( + status_code=429, + api_key_hint="sk-ant-api-primary", + credential_id="cred-backup", # stale id after env refresh (#79156) + ) + + statuses = {entry.id: entry.last_status for entry in pool.entries()} + assert statuses["cred-primary"] == STATUS_EXHAUSTED + assert statuses["cred-backup"] != STATUS_EXHAUSTED + # Rotation hands the healthy backup (or None if selection prefers next). + if next_entry is not None: + assert next_entry.id == "cred-backup" + assert next_entry.runtime_api_key == "sk-ant-api-backup" + + def test_unmatched_api_key_hint_rotates_without_benching_innocent_key(tmp_path, monkeypatch): """An api_key_hint matching no entry must not quarantine a healthy key. diff --git a/tests/run_agent/test_env_credential_turn_refresh.py b/tests/run_agent/test_env_credential_turn_refresh.py index b310a3d629..078baeb8e6 100644 --- a/tests/run_agent/test_env_credential_turn_refresh.py +++ b/tests/run_agent/test_env_credential_turn_refresh.py @@ -112,6 +112,43 @@ class TestLeavesNonEnvStateAlone: assert agent._try_refresh_env_client_credentials() is False assert agent.api_key == "sk-rotated-pool-entry" + def test_first_look_does_not_stomp_rotated_pool_key(self, env): + """#79156: first look with a pool-rotated key must seed the baseline, + not rewrite api_key back to the env primary.""" + agent = _make_agent(api_key="sk-backup") + agent._credential_pool = object() # any non-None pool binding + agent._credential_pool_entry_id = "entry-backup" + env["OPENAI_API_KEY"] = "sk-primary" + + assert agent._try_refresh_env_client_credentials() is False + assert agent.api_key == "sk-backup" + assert agent._credential_pool_entry_id == "entry-backup" + + def test_adopted_key_rebinds_pool_entry_id(self, env, monkeypatch): + """#79156: adopting a new env key must rebind _credential_pool_entry_id + so the next 429 is attributed to the key that actually ran.""" + agent = _make_agent() + env["OPENAI_API_KEY"] = "sk-old" + assert agent._try_refresh_env_client_credentials() is False + agent._credential_pool_entry_id = "stale-rotated-id" + + called = {} + + def fake_sync(a): + called["agent"] = a + a._credential_pool_entry_id = "entry-for-sk-new" + + monkeypatch.setattr( + "agent.agent_runtime_helpers.sync_credential_pool_entry_id", + fake_sync, + ) + + env["OPENAI_API_KEY"] = "sk-new" + assert agent._try_refresh_env_client_credentials() is True + assert agent.api_key == "sk-new" + assert called.get("agent") is agent + assert agent._credential_pool_entry_id == "entry-for-sk-new" + def test_custom_endpoint_wins_over_env_edit(self, env): """A session running on a config/pool custom endpoint (not the registry default, not a previously-seen env value) keeps it."""