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

agent/agent_runtime_helpers.py

@@ -646,6 +646,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:
@@ -657,7 +669,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",
@@ -681,7 +693,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",
@@ -705,7 +717,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",
@@ -781,7 +793,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",

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"
+        )