fix(codex): rotate pool on usage limit 429

This commit is contained in:
Maxim Esipov 2026-05-15 17:19:31 +03:00 committed by Teknium
parent dffb602f37
commit e51d74ab91
2 changed files with 57 additions and 2 deletions

View file

@ -5135,7 +5135,7 @@ class AIAgent:
if isinstance(body, dict): if isinstance(body, dict):
payload = body.get("error") if isinstance(body.get("error"), dict) else body payload = body.get("error") if isinstance(body.get("error"), dict) else body
if isinstance(payload, dict): if isinstance(payload, dict):
reason = payload.get("code") or payload.get("error") reason = payload.get("code") or payload.get("type") or payload.get("error")
if isinstance(reason, str) and reason.strip(): if isinstance(reason, str) and reason.strip():
context["reason"] = reason.strip() context["reason"] = reason.strip()
message = payload.get("message") or payload.get("error_description") message = payload.get("message") or payload.get("error_description")
@ -7583,7 +7583,15 @@ class AIAgent:
return False, has_retried_429 return False, has_retried_429
if effective_reason == FailoverReason.rate_limit: if effective_reason == FailoverReason.rate_limit:
if not has_retried_429: usage_limit_reached = False
if error_context:
context_reason = str(error_context.get("reason") or "").lower()
context_message = str(error_context.get("message") or "").lower()
usage_limit_reached = (
"usage_limit_reached" in context_reason
or "usage limit has been reached" in context_message
)
if not has_retried_429 and not usage_limit_reached:
return False, True return False, True
rotate_status = status_code if status_code is not None else 429 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)

View file

@ -3746,6 +3746,37 @@ class TestCredentialPoolRecovery:
assert retry_same is False assert retry_same is False
agent._swap_credential.assert_called_once_with(next_entry) agent._swap_credential.assert_called_once_with(next_entry)
def test_recover_with_pool_rotates_usage_limit_429_immediately(self, agent):
next_entry = SimpleNamespace(label="secondary")
captured = {}
class _Pool:
def current(self):
return SimpleNamespace(label="primary")
def mark_exhausted_and_rotate(self, *, status_code, error_context=None):
captured["status_code"] = status_code
captured["error_context"] = error_context
return next_entry
agent._credential_pool = _Pool()
agent._swap_credential = MagicMock()
recovered, retry_same = agent._recover_with_credential_pool(
status_code=429,
has_retried_429=False,
error_context={
"reason": "usage_limit_reached",
"message": "The usage limit has been reached",
},
)
assert recovered is True
assert retry_same is False
assert captured["status_code"] == 429
assert captured["error_context"]["reason"] == "usage_limit_reached"
agent._swap_credential.assert_called_once_with(next_entry)
def test_recover_with_pool_refreshes_on_401(self, agent): def test_recover_with_pool_refreshes_on_401(self, agent):
"""401 with successful refresh should swap to refreshed credential.""" """401 with successful refresh should swap to refreshed credential."""
@ -3832,6 +3863,22 @@ class TestCredentialPoolRecovery:
assert context["message"] == "Weekly credits exhausted." assert context["message"] == "Weekly credits exhausted."
assert context["reset_at"] == "2026-04-12T10:30:00Z" assert context["reset_at"] == "2026-04-12T10:30:00Z"
def test_extract_api_error_context_uses_type_as_reason(self, agent):
error = SimpleNamespace(
body={
"error": {
"type": "usage_limit_reached",
"message": "The usage limit has been reached",
}
},
response=SimpleNamespace(headers={}),
)
context = agent._extract_api_error_context(error)
assert context["reason"] == "usage_limit_reached"
assert context["message"] == "The usage limit has been reached"
def test_recover_with_pool_passes_error_context_on_rotated_429(self, agent): def test_recover_with_pool_passes_error_context_on_rotated_429(self, agent):
next_entry = SimpleNamespace(label="secondary") next_entry = SimpleNamespace(label="secondary")
captured = {} captured = {}