diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 263cf1563a27..29f48eeaac3c 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -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", diff --git a/tests/agent/test_credential_pool_routing.py b/tests/agent/test_credential_pool_routing.py index a23c2280096b..95b07a7039d9 100644 --- a/tests/agent/test_credential_pool_routing.py +++ b/tests/agent/test_credential_pool_routing.py @@ -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" + )