fix(providers): preserve quota cache on refresh failure

This commit is contained in:
starship-s
2026-05-12 21:16:34 -06:00
parent a166625e02
commit c562ce2e8c
3 changed files with 62 additions and 3 deletions
+2
View File
@@ -2333,6 +2333,8 @@ def invalidate_credential_pool_cache(provider_id: str):
_CREDENTIAL_POOL_CACHE.pop(provider_id, None)
_CREDENTIAL_POOL_CACHE.pop(_resolve_provider_alias(provider_id), None)
try:
# api.providers imports from api.config; keep this lazy to avoid
# import-cycle/module-initialization issues.
from api.providers import invalidate_account_usage_status_cache
invalidate_account_usage_status_cache(provider_id)
+12 -3
View File
@@ -1176,9 +1176,18 @@ def invalidate_account_usage_status_cache(provider_id: str | None = None) -> Non
_account_usage_status_cache.pop(key, None)
def _set_cached_account_usage(cache_key: tuple[str, str, str], snapshot: Any) -> None:
def _set_cached_account_usage(
cache_key: tuple[str, str, str],
snapshot: Any,
*,
preserve_non_none: bool = False,
) -> None:
now = time.monotonic()
with _account_usage_status_cache_lock:
if preserve_non_none and snapshot is None:
cached = _account_usage_status_cache.get(cache_key)
if cached is not None and cached[1] is not None:
return
_account_usage_status_cache[cache_key] = (now, snapshot)
expired = [
key for key, (fetched_at, _snapshot) in _account_usage_status_cache.items()
@@ -1269,11 +1278,11 @@ def _fetch_account_usage_with_profile_context(provider: str, *, refresh: bool =
home,
api_key=api_key,
)
_set_cached_account_usage(cache_key, snapshot)
_set_cached_account_usage(cache_key, snapshot, preserve_non_none=refresh)
return snapshot
except Exception:
logger.debug("Failed to fetch account usage for %s", provider, exc_info=True)
_set_cached_account_usage(cache_key, None)
_set_cached_account_usage(cache_key, None, preserve_non_none=refresh)
return None
+48
View File
@@ -886,6 +886,54 @@ def test_account_usage_profile_fetch_uses_short_lived_cache(monkeypatch, tmp_pat
]
def test_account_usage_forced_refresh_failure_preserves_warm_snapshot(monkeypatch, tmp_path):
"""A failed manual refresh should not discard the last usable account snapshot."""
import api.providers as providers
monkeypatch.setattr(profiles, "get_active_hermes_home", lambda: tmp_path)
old_cfg, old_mtime = _with_config(model={"provider": "openai-codex"})
providers._account_usage_status_cache.clear()
calls = []
good_snapshot = SimpleNamespace(
provider="openai-codex",
source="usage_api_pool",
title="Account limits",
plan=None,
windows=(),
details=(),
available=True,
unavailable_reason=None,
fetched_at=datetime(2030, 3, 17, 12, 30, tzinfo=timezone.utc),
pool={"total_credentials": 1, "credentials": []},
)
def fake_fetch(provider, home, api_key=None):
calls.append((provider, str(home), api_key))
return good_snapshot if len(calls) == 1 else None
monkeypatch.setattr(providers, "_agent_fetch_account_usage_for_home", fake_fetch)
try:
first = providers._fetch_account_usage_with_profile_context("openai-codex")
refreshed = providers._fetch_account_usage_with_profile_context(
"openai-codex",
refresh=True,
)
after_refresh_failure = providers._fetch_account_usage_with_profile_context(
"openai-codex",
)
finally:
providers._account_usage_status_cache.clear()
_restore_config(old_cfg, old_mtime)
assert first is good_snapshot
assert refreshed is None
assert after_refresh_failure is good_snapshot
assert calls == [
("openai-codex", str(tmp_path), None),
("openai-codex", str(tmp_path), None),
]
def test_account_usage_profile_cache_invalidates_with_credential_pool_cache(monkeypatch, tmp_path):
"""Credential-pool invalidation should also clear pooled account usage."""
import api.providers as providers