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
This commit is contained in:
686f6c61
2026-08-05 09:39:02 +02:00
committed by Teknium
parent 0b33ee88e4
commit bf7c716648
4 changed files with 147 additions and 0 deletions
+31
View File
@@ -2045,6 +2045,37 @@ class CredentialPool:
(e for e in self._entries if e.id == credential_id), (e for e in self._entries if e.id == credential_id),
None, 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: if entry is None and api_key_hint:
# Prefer the specific entry whose API key matches the one that # Prefer the specific entry whose API key matches the one that
# actually failed. When this pool was freshly loaded from disk # actually failed. When this pool was freshly loaded from disk
+24
View File
@@ -5611,6 +5611,17 @@ class AIAgent:
adopt = current_base == default_base and not ( adopt = current_base == default_base and not (
base_url == current_base and api_key == self.api_key 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: else:
# Env unchanged → no-op; any drift from self.* is rotation/ # Env unchanged → no-op; any drift from self.* is rotation/
# failover or config precedence — leave it alone. An edit is # failover or config precedence — leave it alone. An edit is
@@ -5654,6 +5665,19 @@ class AIAgent:
self._client_kwargs.update(prior_client_kwargs) self._client_kwargs.update(prior_client_kwargs)
return False 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 self._env_creds_seen = resolved
logger.info( logger.info(
"Applied updated .env credentials for %s: endpoint %s", "Applied updated .env credentials for %s: endpoint %s",
+55
View File
@@ -143,6 +143,61 @@ def test_billing_rotation_marks_all_entries_sharing_failed_key(tmp_path, monkeyp
assert statuses["cred-model-config"] == STATUS_EXHAUSTED 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): 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. """An api_key_hint matching no entry must not quarantine a healthy key.
@@ -112,6 +112,43 @@ class TestLeavesNonEnvStateAlone:
assert agent._try_refresh_env_client_credentials() is False assert agent._try_refresh_env_client_credentials() is False
assert agent.api_key == "sk-rotated-pool-entry" 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): def test_custom_endpoint_wins_over_env_edit(self, env):
"""A session running on a config/pool custom endpoint (not the """A session running on a config/pool custom endpoint (not the
registry default, not a previously-seen env value) keeps it.""" registry default, not a previously-seen env value) keeps it."""