mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-10-20 05:15:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

fix(agent): pass api_key_hint to mark_exhausted_and_rotate in credential pool recovery

Browse source 
recover_with_credential_pool() called mark_exhausted_and_rotate() without
api_key_hint, causing it to fall back to current() or _select_unlocked().
When a prior rotation left current() as None, _select_unlocked() returned
the NEXT (healthy) entry instead of the one that actually failed — marking
the wrong credential as exhausted (#43747).

Extract the current API key from agent.api_key (or pool.current().runtime_api_key
as fallback) and pass it as api_key_hint to all 4 call sites.
This commit is contained in:
liuhao1024 2026-06-11 04:11:18 +08:00 • committed by Teknium
parent 3f9944bad9
commit 702f5f10ad
2 changed files with 62 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

26
agent/agent_runtime_helpers.py
View file

@ -922,6 +922,18 @@ def recover_with_credential_pool(
)
return False, has_retried_429
# Capture the current API key before any rotation — needed to
# identify which credential actually failed when
# mark_exhausted_and_rotate is called. Without this hint the
# pool falls back to current() or _select_unlocked(), which may
# return the NEXT (healthy) entry after a prior rotation, marking
# the wrong credential as exhausted (#43747).
_api_key_hint = getattr(agent, "api_key", None) or None
if not _api_key_hint:
_cur = pool.current()
if _cur:
_api_key_hint = getattr(_cur, "runtime_api_key", None)
effective_reason = classified_reason
if effective_reason is None:
if status_code == 402:
@ -952,13 +964,13 @@ def recover_with_credential_pool(
if effective_reason == FailoverReason.billing:
rotate_status = status_code if status_code is not None else 402
# Runtime credentials can be resolved by a separate pool instance,
# leaving this recovery pool without ``current_id``. Match the key
# that actually failed instead of quarantining a different account.
next_entry = pool.mark_exhausted_and_rotate(
status_code=rotate_status,
error_context=error_context,
# Runtime credentials can be resolved by a separate pool instance,
# leaving this recovery pool without ``current_id``. Match the key
# that actually failed instead of quarantining a different account.
api_key_hint=getattr(agent, "api_key", None),
api_key_hint=_api_key_hint,
)
if next_entry is not None:
_ra().logger.info(
@ -983,7 +995,7 @@ def recover_with_credential_pool(
current_last_status,
)
rotate_status = status_code if status_code is not None else 429
next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context)
next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context, api_key_hint=_api_key_hint)
if next_entry is not None:
_ra().logger.info(
"Credential %s (rate limit, pre-exhausted) — rotated to pool entry %s",
@ -1007,7 +1019,7 @@ def recover_with_credential_pool(
if not has_retried_429 and not usage_limit_reached:
return False, True
rotate_status = status_code if status_code is not None else 429
next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context)
next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context, api_key_hint=_api_key_hint)
if next_entry is not None:
_ra().logger.info(
"Credential %s (rate limit) — rotated to pool entry %s",
@ -1107,7 +1119,7 @@ def recover_with_credential_pool(
# Refresh failed — rotate to next credential instead of giving up.
# The failed entry is already marked exhausted by try_refresh_current().
rotate_status = status_code if status_code is not None else 401
next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context)
next_entry = pool.mark_exhausted_and_rotate(status_code=rotate_status, error_context=error_context, api_key_hint=_api_key_hint)
if next_entry is not None:
_ra().logger.info(
"Credential %s (auth refresh failed) — rotated to pool entry %s",

50
tests/agent/test_credential_pool_routing.py
View file

@ -162,7 +162,7 @@ class TestPoolRotationCycle:
# mark_exhausted_and_rotate returns next entry until exhausted
self._rotation_index = 0
def rotate(status_code=None, error_context=None, api_key_hint=None):
def rotate(status_code=None, error_context=None, **_kwargs):
self._rotation_index += 1
if self._rotation_index < pool_entries:
return entries[self._rotation_index]
@ -173,6 +173,9 @@ class TestPoolRotationCycle:
agent._credential_pool = pool
agent._swap_credential = MagicMock()
agent.log_prefix = ""
agent.api_key = "test-api-key"
agent.provider = "test-provider"
pool.provider = "test-provider"
return agent, pool, entries
@ -194,7 +197,7 @@ class TestPoolRotationCycle:
)
assert recovered is True
assert has_retried is False # reset after rotation
pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=429, error_context=None)
pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=429, error_context=None, api_key_hint="test-api-key")
agent._swap_credential.assert_called_once_with(entries[1])
def test_pool_exhaustion_returns_false(self):
@ -220,11 +223,7 @@ class TestPoolRotationCycle:
)
assert recovered is True
assert has_retried is False
pool.mark_exhausted_and_rotate.assert_called_once_with(
status_code=402,
error_context=None,
api_key_hint=None,
)
pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=402, error_context=None, api_key_hint="test-api-key")
def test_no_pool_returns_false(self):
"""No pool should return (False, unchanged)."""
@ -239,3 +238,40 @@ class TestPoolRotationCycle:
)
assert recovered is False
assert has_retried is False
def test_api_key_hint_from_pool_current_when_agent_key_missing(self):
"""api_key_hint should fall back to pool.current().runtime_api_key
when agent.api_key is not set (#43747)."""
from run_agent import AIAgent
with patch.object(AIAgent, "__init__", lambda self, **kw: None):
agent = AIAgent()
e0 = MagicMock(name="entry_0")
e0.id = "cred-0"
e1 = MagicMock(name="entry_1")
e1.id = "cred-1"
pool = MagicMock()
pool.has_credentials.return_value = True
pool.provider = "test-provider"
agent.provider = "test-provider"
# current entry has a runtime_api_key
cur_entry = MagicMock()
cur_entry.runtime_api_key = "pool-current-key"
pool.current.return_value = cur_entry
pool.mark_exhausted_and_rotate.return_value = e1
agent._credential_pool = pool
agent._swap_credential = MagicMock()
agent.log_prefix = ""
# No agent.api_key set — should fall back to pool.current().runtime_api_key
recovered, has_retried = agent._recover_with_credential_pool(
status_code=402, has_retried_429=False
)
assert recovered is True
pool.mark_exhausted_and_rotate.assert_called_once_with(
status_code=402, error_context=None, api_key_hint="pool-current-key"
)