From c562ce2e8c989398b52081e9a26d358e05e7f1bb Mon Sep 17 00:00:00 2001
From: starship-s <45587122+starship-s@users.noreply.github.com>
Date: Tue, 12 May 2026 21:16:34 -0600
Subject: [PATCH] fix(providers): preserve quota cache on refresh failure

---
 api/config.py                       |  2 ++
 api/providers.py                    | 15 +++++++--
 tests/test_provider_quota_status.py | 48 +++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/api/config.py b/api/config.py
index a500900b..68841eec 100644
--- a/api/config.py
+++ b/api/config.py
@@ -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)
diff --git a/api/providers.py b/api/providers.py
index 49b0b6a0..0cd869fa 100644
--- a/api/providers.py
+++ b/api/providers.py
@@ -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
 
 
diff --git a/tests/test_provider_quota_status.py b/tests/test_provider_quota_status.py
index f407e27c..ae789e13 100644
--- a/tests/test_provider_quota_status.py
+++ b/tests/test_provider_quota_status.py
@@ -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