From d1eb08fcf3a22c9345a1adeaa63812a2cde7f32d Mon Sep 17 00:00:00 2001 From: kshitij Date: Tue, 4 Aug 2026 15:14:15 +0530 Subject: [PATCH] fix: thread sole_credential into next_available_at sibling site MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit next_available_at() was computing the full 1-hour TTL for a sole credential on a 429, contradicting the 60s cooldown in _available_entries. The fallback restore gate (agent_runtime_helpers) uses next_available_at to decide when to switch back from fallback to primary — so the agent stayed on fallback for an hour instead of ~60s. Add sole_credential computation in next_available_at mirroring _available_entries, and a test verifying the short cooldown propagates. --- agent/credential_pool.py | 9 ++++++- .../test_credential_pool_sole_cooldown.py | 24 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/agent/credential_pool.py b/agent/credential_pool.py index 6d2cd42c3c0d6..3466464424394 100644 --- a/agent/credential_pool.py +++ b/agent/credential_pool.py @@ -665,11 +665,18 @@ class CredentialPool: available, _pending = self._available_entries() if available: return None + # Mirror _available_entries: if the pool has no other credential + # to rotate to, the sole entry's transient throttle cools down in + # seconds — next_available_at must report that shorter window too, + # or the fallback restore gate waits an hour for a 60s cooldown. + sole_credential = sum( + 1 for e in self._entries if e.last_status != STATUS_DEAD + ) <= 1 candidates: List[float] = [] for entry in self._entries: if entry.last_status != STATUS_EXHAUSTED: continue - until = _exhausted_until(entry) + until = _exhausted_until(entry, sole_credential=sole_credential) if until is not None: candidates.append(until) return min(candidates) if candidates else None diff --git a/tests/agent/test_credential_pool_sole_cooldown.py b/tests/agent/test_credential_pool_sole_cooldown.py index 37b134bc2ae58..ea933ba93a5cd 100644 --- a/tests/agent/test_credential_pool_sole_cooldown.py +++ b/tests/agent/test_credential_pool_sole_cooldown.py @@ -76,6 +76,30 @@ def test_sole_credential_402_keeps_full_bench(tmp_path, monkeypatch): assert pool.select() is None +def test_sole_credential_next_available_at_uses_short_cooldown(tmp_path, monkeypatch): + """next_available_at must also honour the sole-credential short cooldown. + + Without this, the fallback restore gate in agent_runtime_helpers waits an + hour for a 60s cooldown, keeping the agent on a fallback provider far + longer than necessary. + """ + from agent.credential_pool import EXHAUSTED_TTL_SOLE_CREDENTIAL_SECONDS + + pool = _load(tmp_path, monkeypatch, [_entry(429, age_seconds=10)]) + next_at = pool.next_available_at() + assert next_at is not None + # Should be ~60s from exhaustion, not ~3600s. The entry was exhausted 10s + # ago, so the remaining wait is ~50s. + remaining = next_at - time.time() + assert remaining < EXHAUSTED_TTL_SOLE_CREDENTIAL_SECONDS, ( + f"next_available_at returned {remaining:.0f}s remaining — expected < " + f"{EXHAUSTED_TTL_SOLE_CREDENTIAL_SECONDS}s (sole-credential cooldown)" + ) + assert remaining < 300, ( + f"next_available_at returned {remaining:.0f}s — should be seconds, not hours" + ) + + def test_multi_key_429_keeps_full_bench(tmp_path, monkeypatch): """With more than one non-DEAD entry there IS something to rotate to, so the short cooldown must not kick in — both recently-throttled keys stay benched."""