From 9c159d601e4837600f8276175107ec2e0ffe31ec Mon Sep 17 00:00:00 2001 From: liuhao1024 Date: Thu, 11 Jun 2026 04:11:18 +0800 Subject: [PATCH] fix(agent): pass api_key_hint to mark_exhausted_and_rotate in credential pool recovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- agent/agent_runtime_helpers.py | 20 +++++++-- tests/agent/test_credential_pool_routing.py | 46 +++++++++++++++++++-- 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index daffc025d9bd..cc5f891721ad 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -586,6 +586,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: @@ -597,7 +609,7 @@ def recover_with_credential_pool( if effective_reason == FailoverReason.billing: rotate_status = status_code if status_code is not None else 402 - 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 (billing) — rotated to pool entry %s", @@ -621,7 +633,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", @@ -645,7 +657,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", @@ -708,7 +720,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", diff --git a/tests/agent/test_credential_pool_routing.py b/tests/agent/test_credential_pool_routing.py index 8477fdb646a8..c6040a966c46 100644 --- a/tests/agent/test_credential_pool_routing.py +++ b/tests/agent/test_credential_pool_routing.py @@ -159,7 +159,7 @@ def _make_agent_with_pool(self, pool_entries=3): # mark_exhausted_and_rotate returns next entry until exhausted self._rotation_index = 0 - def rotate(status_code=None, error_context=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] @@ -170,6 +170,9 @@ def rotate(status_code=None, error_context=None): 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 @@ -191,7 +194,7 @@ def test_second_429_rotates_to_next(self): ) 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): @@ -217,7 +220,7 @@ def test_402_immediate_rotation(self): ) assert recovered is True assert has_retried is False - pool.mark_exhausted_and_rotate.assert_called_once_with(status_code=402, error_context=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).""" @@ -232,3 +235,40 @@ def test_no_pool_returns_false(self): ) 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" + )