diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 606edbad7836..a418088916a9 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -1056,7 +1056,7 @@ def recover_with_credential_pool( # Subscription/entitlement 403s look like auth failures on the wire # but refresh cannot fix them — the OAuth token is already valid, # the account simply lacks the entitlement. Without this guard, - # ``try_refresh_current()`` keeps minting fresh tokens against the + # the refresh path keeps minting fresh tokens against the # same unsubscribed account and the main agent loop spins re-issuing # the same 403 until the user Ctrl+C's. # @@ -1109,9 +1109,13 @@ def recover_with_credential_pool( agent.provider or "provider", ) return False, has_retried_429 - refreshed = pool.try_refresh_current() + # Refresh the entry that supplied the failing key, not current(): + # the shared pointer can reference a different, healthy entry, and + # refreshing it would consume that entry's single-use refresh token + # (or mark it exhausted on failure) for a failure it never had. + refreshed = pool.try_refresh_matching(api_key_hint=_api_key_hint) if refreshed is not None: - # ``try_refresh_current()`` re-mints a fresh OAuth token and reports + # ``try_refresh_matching()`` re-mints a fresh OAuth token and reports # success even when the upstream keeps rejecting it — a single-entry # pool (common for OAuth/Max subscribers) has nothing to rotate to, # so a bare "refreshed → retry" loop spins forever on the same dead @@ -1139,7 +1143,7 @@ def recover_with_credential_pool( agent._swap_credential(refreshed) return True, has_retried_429 # Refresh failed — rotate to next credential instead of giving up. - # The failed entry is already marked exhausted by try_refresh_current(). + # The failed entry is already marked exhausted by the refresh attempt. rotate_status = status_code if status_code is not None else 401 next_entry = pool.mark_exhausted_and_rotate( status_code=rotate_status, diff --git a/tests/agent/test_credential_pool_routing.py b/tests/agent/test_credential_pool_routing.py index 54b63f90c008..8f093f86b0e1 100644 --- a/tests/agent/test_credential_pool_routing.py +++ b/tests/agent/test_credential_pool_routing.py @@ -489,3 +489,35 @@ class TestFailureAttribution: assert statuses["cred-0"] != "exhausted" swapped = agent._swap_credential.call_args[0][0] assert swapped.id == "cred-0" + + def test_auth_refresh_targets_failing_key_not_pointer(self, tmp_path, monkeypatch): + """The auth path must refresh the entry that supplied the failing key, + not current(). With current() pointing at healthy A while key B failed, + try_refresh_current() force-refreshes A — for non-OAuth entries a + forced refresh marks the entry exhausted outright — so healthy A dies, + the hinted rotation then exhausts B, and the pool has nothing left.""" + pool = self._make_pool( + tmp_path, monkeypatch, + [self._entry(0, "key-a"), self._entry(1, "key-b")], + ) + # Point the shared cursor at the healthy entry, as a concurrent + # turn's select() would. + selected = pool.select() + assert selected.id == "cred-0" + assert pool.current().id == "cred-0" + + agent = self._agent(pool, failing_key="key-b") + agent._is_entitlement_failure = MagicMock(return_value=False) + + from agent.agent_runtime_helpers import recover_with_credential_pool + + recovered, _ = recover_with_credential_pool( + agent, status_code=401, has_retried_429=False + ) + + assert recovered is True + statuses = self._statuses(pool) + assert statuses["cred-1"] == "exhausted" + assert statuses["cred-0"] != "exhausted" + swapped = agent._swap_credential.call_args[0][0] + assert swapped.id == "cred-0" diff --git a/tests/run_agent/test_codex_xai_oauth_recovery.py b/tests/run_agent/test_codex_xai_oauth_recovery.py index cb926a09affe..f82b216357bb 100644 --- a/tests/run_agent/test_codex_xai_oauth_recovery.py +++ b/tests/run_agent/test_codex_xai_oauth_recovery.py @@ -706,7 +706,7 @@ def test_recover_with_credential_pool_skips_refresh_on_entitlement_403(): refresh_calls = {"n": 0} class _FakePool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): refresh_calls["n"] += 1 return MagicMock(id="should_not_be_called") @@ -756,7 +756,7 @@ def test_recover_with_credential_pool_rotates_on_xai_spending_limit_403(): class _FakePool: provider = "xai-oauth" - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): refresh_calls["n"] += 1 return MagicMock(id="should_not_be_called") @@ -816,7 +816,7 @@ def test_recover_with_credential_pool_skips_refresh_on_bare_403_for_xai_oauth(): refresh_calls = {"n": 0} class _FakePool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): refresh_calls["n"] += 1 return MagicMock(id="should_not_be_called") @@ -857,7 +857,7 @@ def test_recover_with_credential_pool_still_refreshes_genuine_auth_failure(): refresh_calls = {"n": 0} class _FakePool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): refresh_calls["n"] += 1 # Return a fake refreshed entry — semantically "refresh worked" entry = MagicMock() @@ -1038,7 +1038,7 @@ def test_recover_with_credential_pool_refreshes_on_xai_bad_credentials_403(): refresh_calls = {"n": 0} class _FakePool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): refresh_calls["n"] += 1 entry = MagicMock() entry.id = "entry_refreshed_after_stale" @@ -1094,7 +1094,7 @@ def test_recover_with_credential_pool_still_blocks_real_entitlement(): refresh_calls = {"n": 0} class _FakePool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): refresh_calls["n"] += 1 return MagicMock(id="should_not_be_called") diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index f4edc65896bb..98f5a779ca02 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -6661,7 +6661,7 @@ class TestCredentialPoolRecovery: refreshed_entry = SimpleNamespace(label="refreshed-primary", id="abc") class _Pool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): return refreshed_entry agent._credential_pool = _Pool() @@ -6679,7 +6679,7 @@ class TestCredentialPoolRecovery: """Repeated same-entry auth refreshes must eventually fall through. A single-entry OAuth pool re-mints a fresh token on every 401, so - ``try_refresh_current()`` reports success forever. The cap (#26080) + ``try_refresh_matching()`` reports success forever. The cap (#26080) must let the third consecutive same-entry refresh fall through (return not-recovered) so the fallback chain can activate instead of looping on the same dead credential. @@ -6687,7 +6687,7 @@ class TestCredentialPoolRecovery: refreshed_entry = SimpleNamespace(label="primary", id="abc") class _Pool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): return refreshed_entry agent._credential_pool = _Pool() @@ -6711,7 +6711,7 @@ class TestCredentialPoolRecovery: sequence = [entry_a, entry_a, entry_b, entry_b] class _Pool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): return sequence.pop(0) agent._credential_pool = _Pool() @@ -6733,7 +6733,7 @@ class TestCredentialPoolRecovery: next_entry = SimpleNamespace(label="secondary", id="def") class _Pool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): return None # refresh failed def mark_exhausted_and_rotate( @@ -6760,7 +6760,7 @@ class TestCredentialPoolRecovery: """401 with failed refresh and no other credentials returns not recovered.""" class _Pool: - def try_refresh_current(self): + def try_refresh_matching(self, api_key_hint=None): return None def mark_exhausted_and_rotate(