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:
@@ -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
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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."""
|
||||||
|
|||||||
Reference in New Issue
Block a user