From 1904eaed6bf531cd6485ba1bc1a5ffb550bfb441 Mon Sep 17 00:00:00 2001 From: starship-s <45587122+starship-s@users.noreply.github.com> Date: Tue, 12 May 2026 20:05:20 -0600 Subject: [PATCH] fix(providers): show fallback pool cooldown times --- api/providers.py | 28 ++++++++-- tests/test_provider_quota_status.py | 80 ++++++++++++++++++++++++++++- 2 files changed, 103 insertions(+), 5 deletions(-) diff --git a/api/providers.py b/api/providers.py index 708759d9..e4f316f3 100644 --- a/api/providers.py +++ b/api/providers.py @@ -19,7 +19,7 @@ import threading import time import urllib.error import urllib.request -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace from typing import Any @@ -141,7 +141,7 @@ _ACCOUNT_USAGE_SUBPROCESS_CODE = r""" import base64 import json import sys -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from urllib import request as urllib_request @@ -384,13 +384,33 @@ def _safe_unavailable_reason(reason): return text[:180] +def _entry_exhausted_ttl_seconds(error_code): + code = str(error_code or "").strip() + if code == "401": + return 5 * 60 + return 60 * 60 + + +def _entry_pool_exhausted_until(entry): + if str(_entry_value(entry, "last_status") or "").strip().lower() != "exhausted": + return None + reset_at = _parse_dt(getattr(entry, "last_error_reset_at", None)) + if reset_at is not None: + return reset_at + status_at = _parse_dt(getattr(entry, "last_status_at", None)) + if status_at is None: + return None + return status_at + timedelta(seconds=_entry_exhausted_ttl_seconds(_entry_value(entry, "last_error_code"))) + + def _entry_is_pool_exhausted(entry): - return str(_entry_value(entry, "last_status") or "").strip().lower() == "exhausted" + exhausted_until = _entry_pool_exhausted_until(entry) + return exhausted_until is not None and datetime.now(timezone.utc) < exhausted_until def _entry_pool_exhausted_reason(entry): code = _entry_value(entry, "last_error_code") - reset_at = _iso(_parse_dt(getattr(entry, "last_error_reset_at", None))) + reset_at = _iso(_entry_pool_exhausted_until(entry)) reason = "Credential pool marked this credential exhausted" if code: reason += " after provider status " + code diff --git a/tests/test_provider_quota_status.py b/tests/test_provider_quota_status.py index 5ca97db8..052887d9 100644 --- a/tests/test_provider_quota_status.py +++ b/tests/test_provider_quota_status.py @@ -328,6 +328,7 @@ def test_codex_account_usage_subprocess_reports_read_only_credential_pool(monkey runtime_api_key=exhausted_token, runtime_base_url="https://chatgpt.com/backend-api/codex", last_status="exhausted", + last_status_at=1_900_000_000, ), ] @@ -440,7 +441,7 @@ def test_codex_account_usage_subprocess_reports_read_only_credential_pool(monkey "plan": None, "windows": [], "details": [], - "unavailable_reason": "Credential pool marked this credential exhausted.", + "unavailable_reason": "Credential pool marked this credential exhausted; retry after 2030-03-17T18:46:40Z.", "fetched_at": None, }, ], @@ -449,6 +450,83 @@ def test_codex_account_usage_subprocess_reports_read_only_credential_pool(monkey assert exhausted_token not in output +def test_codex_account_usage_subprocess_retries_expired_pool_exhaustion(monkeypatch, capsys): + """Expired pool cooldowns should be probed instead of shown as still exhausted.""" + import api.providers as providers + + def b64url(payload: bytes) -> str: + return base64.urlsafe_b64encode(payload).rstrip(b"=").decode("ascii") + + token = ".".join(( + b64url(b'{"alg":"none","typ":"JWT"}'), + b64url(json.dumps({ + "https://api.openai.com/auth": { + "chatgpt_account_id": "acct-expired", + }, + }).encode("utf-8")), + b64url(b"signature"), + )) + seen = [] + + agent_mod = types.ModuleType("agent") + agent_mod.__path__ = [] + account_usage_mod = types.ModuleType("agent.account_usage") + credential_pool_mod = types.ModuleType("agent.credential_pool") + + def fake_fetch_account_usage(provider, *, base_url=None, api_key=None): + return None + + class FakePool: + def entries(self): + return [ + SimpleNamespace( + label="Expired cooldown", + runtime_api_key=token, + runtime_base_url="https://chatgpt.com/backend-api/codex", + last_status="exhausted", + last_status_at=1, + last_error_code=None, + last_error_reset_at=None, + ), + ] + + def select(self): + raise AssertionError("quota display must not rotate credential_pool selection") + + def fake_load_pool(provider): + return FakePool() + + def fake_urlopen(req, timeout): + seen.append(req.full_url) + payload = { + "plan_type": "team", + "rate_limit": { + "primary_window": {"used_percent": 10, "reset_at": "2030-03-17T17:30:00Z"}, + }, + } + return _FakeResponse(json.dumps(payload).encode("utf-8")) + + account_usage_mod.fetch_account_usage = fake_fetch_account_usage + credential_pool_mod.load_pool = fake_load_pool + monkeypatch.setitem(sys.modules, "agent", agent_mod) + monkeypatch.setitem(sys.modules, "agent.account_usage", account_usage_mod) + monkeypatch.setitem(sys.modules, "agent.credential_pool", credential_pool_mod) + monkeypatch.setattr(providers.urllib.request, "urlopen", fake_urlopen) + monkeypatch.setattr(sys, "argv", ["quota-probe", "openai-codex", ""]) + + exec(providers._ACCOUNT_USAGE_SUBPROCESS_CODE, {"__name__": "__main__"}) + + output = capsys.readouterr().out.strip() + snapshot = json.loads(output) + + assert seen == ["https://chatgpt.com/backend-api/wham/usage"] + assert snapshot["pool"]["queried_credentials"] == 1 + assert snapshot["pool"]["exhausted_credentials"] == 0 + assert snapshot["pool"]["credentials"][0]["status"] == "available" + assert snapshot["pool"]["credentials"][0]["unavailable_reason"] is None + assert token not in output + + def test_codex_account_usage_subprocess_sanitizes_pool_entry_errors(monkeypatch, capsys): """Pool per-entry failures must not leak bearer/JWT-like exception text."""