fix: thread sole_credential into next_available_at sibling site

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.
This commit is contained in:
kshitij 2026-08-04 15:14:15 +05:30 committed by kshitij
parent dcd7504349
commit d1eb08fcf3
2 changed files with 32 additions and 1 deletions

View File

@ -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

View File

@ -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."""