fix(proxy): own cancellable usage refresh sessions

app/modules/proxy/_service/rate_limit.py

@@ -7,6 +7,7 @@
 from app.core import usage as usage_core
 from app.core.usage.types import UsageWindowRow
 from app.db.models import Account, UsageHistory
+from app.modules.accounts.background_repository import BackgroundAccountsRepository
 from app.modules.proxy.helpers import (
     _credits_headers,
     _credits_snapshot,
@@ -26,6 +27,7 @@
     RateLimitWindowSnapshotData,
 )
 from app.modules.usage.additional_quota_keys import get_additional_display_label_for_quota_key
+from app.modules.usage.background_repository import BackgroundAdditionalUsageRepository, BackgroundUsageRepository
 from app.modules.usage.mappers import usage_history_to_window_row
 from app.modules.usage.updater import UsageUpdater
 
@@ -130,11 +132,17 @@ async def get_rate_limit_payload(self) -> RateLimitStatusPayloadData:
         proxy = cast(_RateLimitServiceProtocol, self)
         async with proxy._repo_factory() as repos:
             accounts = await repos.accounts.list_accounts()
-            await self._refresh_usage(repos, accounts)
+            latest_usage = await repos.usage.latest_by_account(window="primary")
+
+        # Do not hold the request session while an owned refresh checks out its
+        # background session; pool_size=1 must still make progress.
+        await self._refresh_usage(accounts, latest_usage)
+
+        async with proxy._repo_factory() as repos:
+            accounts = await repos.accounts.list_accounts()
             selected_accounts = _select_accounts_for_limits(accounts)
             if not selected_accounts:
                 return RateLimitStatusPayloadData(plan_type="guest")
-
             account_map = {account.id: account for account in selected_accounts}
             primary_rows_raw = await self._latest_usage_rows(repos, account_map, "primary")
             secondary_rows_raw = await self._latest_usage_rows(repos, account_map, "secondary")
@@ -176,10 +184,22 @@ async def get_rate_limit_payload(self) -> RateLimitStatusPayloadData:
                 additional_rate_limits=additional_rate_limits,
             )
 
-    async def _refresh_usage(self, repos: ProxyRepositories, accounts: list[Account]) -> None:
-        latest_usage = await repos.usage.latest_by_account(window="primary")
-        updater = UsageUpdater(repos.usage, repos.accounts, repos.additional_usage)
-        await updater.refresh_accounts(accounts, latest_usage)
+    async def _refresh_usage(
+        self,
+        accounts: list[Account],
+        latest_usage: Mapping[str, UsageHistory],
+    ) -> None:
+        updater = UsageUpdater(
+            BackgroundUsageRepository(),
+            BackgroundAccountsRepository(),
+            BackgroundAdditionalUsageRepository(),
+        )
+        await updater.refresh_accounts(
+            accounts,
+            latest_usage,
+            own_singleflight_sessions=True,
+            join_existing=True,
+        )
 
     async def _latest_usage_rows(
         self,

app/modules/usage/updater.py

@@ -276,11 +276,23 @@ async def refresh_accounts(
         latest_usage: Mapping[str, UsageHistory],
         *,
         own_singleflight_sessions: bool = False,
+        join_existing: bool | None = None,
     ) -> bool:
-        """Refresh usage for all accounts. Returns True if usage rows were written."""
+        """Refresh usage for all accounts. Returns True if usage rows were written.
+
+        ``own_singleflight_sessions`` makes each detached singleflight refresh
+        acquire and release its own DB session instead of using this updater's
+        caller-bound repositories. ``join_existing`` controls whether a caller
+        joins an in-flight refresh for the same key (deduplication) or waits
+        and forces a fresh one; it defaults to the historical coupling
+        ``not own_singleflight_sessions`` so existing callers keep their
+        semantics.
+        """
         settings = get_settings()
         if not settings.usage_refresh_enabled:
             return False
+        if join_existing is None:
+            join_existing = not own_singleflight_sessions
 
         refreshed = False
         now = utcnow()
@@ -349,9 +361,9 @@ async def refresh_factory(account: Account = account) -> AccountRefreshResult:
                         own_singleflight_session=own_singleflight_sessions,
                     ),
                     refresh_factory,
-                    join_existing=not own_singleflight_sessions,
+                    join_existing=join_existing,
                 )
-                if not own_singleflight_sessions:
+                if join_existing:
                     await self._sync_account_from_repo(account)
                 refreshed = refreshed or result.usage_written
                 # Only cache when the upstream fetch actually succeeded.
@@ -911,7 +923,9 @@ async def _recover_quota_status_from_usage(
     async def _sync_account_from_repo(self, account: Account) -> None:
         if not self._accounts_repo:
             return
-        stored = await self._accounts_repo.get_by_id(account.id)
+        # Joined owned-session refreshes run in a different session.  A plain
+        # get_by_id() can return the caller session's stale identity-map row.
+        stored = await self._accounts_repo.get_by_id_fresh(account.id)
         if stored is None:
             return
         account.chatgpt_account_id = stored.chatgpt_account_id

tests/unit/test_proxy_rate_limit.py

@@ -17,6 +17,7 @@
 from app.modules.quota_planner.repository import QuotaPlannerRepository
 from app.modules.request_logs.repository import RequestLogsRepository
 from app.modules.usage.repository import AdditionalUsageRepository, UsageRepository
+from app.modules.usage.updater import UsageUpdater
 
 pytestmark = pytest.mark.unit
 
@@ -81,11 +82,20 @@ def __init__(self, repo_factory: ProxyRepoFactory) -> None:
         self._repo_factory = repo_factory
         self.refresh_calls = 0
 
-    async def _refresh_usage(self, repos: ProxyRepositories, accounts: list[Account]) -> None:
-        del repos, accounts
+    async def _refresh_usage(
+        self,
+        accounts: list[Account],
+        latest_usage: dict[str, UsageHistory],
+    ) -> None:
+        del accounts, latest_usage
         self.refresh_calls += 1
 
 
+
+class _RefreshRateLimitService(_RateLimitMixin):
+    def __init__(self, repo_factory: ProxyRepoFactory) -> None:
+        self._repo_factory = repo_factory
+
 def _account(account_id: str, *, plan_type: str) -> Account:
     return Account(
         id=account_id,
@@ -221,6 +231,8 @@ async def test_rate_limit_payload_serializes_usage_reads(monkeypatch: pytest.Mon
 
     assert guard.max_in_flight == 1
     assert guard.calls == [
+        "accounts",
+        "usage:primary",
         "accounts",
         "usage:primary",
         "usage:secondary",
@@ -247,3 +259,32 @@ async def test_rate_limit_payload_serializes_usage_reads(monkeypatch: pytest.Mon
     assert payload.credits.unlimited is False
     assert payload.credits.balance == "8.75"
     assert payload.additional_rate_limits == []
+
+
+@pytest.mark.asyncio
+async def test_rate_limit_usage_refresh_owns_and_joins_singleflight_session(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    base_service, _ = _service_and_guard()
+    service = _RefreshRateLimitService(base_service._repo_factory)
+    account = _account("plus", plan_type="plus")
+    captured: dict[str, object] = {}
+
+    async def capture_refresh(
+        self,
+        accounts: list[Account],
+        latest_usage: dict[str, UsageHistory],
+        **kwargs: object,
+    ) -> bool:
+        del self
+        captured["accounts"] = accounts
+        captured["latest_usage"] = latest_usage
+        captured.update(kwargs)
+        return False
+
+    monkeypatch.setattr(UsageUpdater, "refresh_accounts", capture_refresh)
+
+    await service._refresh_usage([account], {})
+
+    assert captured["own_singleflight_sessions"] is True
+    assert captured["join_existing"] is True

tests/unit/test_usage_updater.py

@@ -82,6 +82,8 @@ async def test_refresh_accounts_owned_singleflight_session_outlives_caller_cance
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     account = _make_account("acc_owned_session", "workspace_owned")
+    stored_account = _make_account("acc_owned_session", "workspace_owned")
+    stored_account.status = AccountStatus.PAUSED
     refresh_started = asyncio.Event()
     allow_refresh_finish = asyncio.Event()
     non_owned_started = asyncio.Event()
@@ -140,6 +142,9 @@ def __init__(self, session) -> None:
         async def get_by_id(self, account_id: str):
             return account if account_id == account.id else None
 
+        async def get_by_id_fresh(self, account_id: str):
+            return stored_account if account_id == account.id else None
+
     @asynccontextmanager
     async def recording_background_session():
         try:
@@ -191,6 +196,7 @@ async def factory() -> usage_updater_module.AccountRefreshResult:
             [account],
             {},
             own_singleflight_sessions=True,
+            join_existing=True,
         )
     )
     await asyncio.wait_for(refresh_started.wait(), timeout=1)
@@ -203,12 +209,87 @@ async def factory() -> usage_updater_module.AccountRefreshResult:
     assert not inner_session_closed.is_set()
     allow_refresh_finish.set()
     await asyncio.wait_for(inner_session_closed.wait(), timeout=1)
+    assert session_was_open_during_refresh == [True]
     allow_non_owned_finish.set()
     await non_owned_task
     await prefixed_non_owned_task
     assert session_was_open_during_refresh == [True]
 
 
+@pytest.mark.parametrize(
+    ("join_existing", "expected_calls"),
+    [(True, 1), (False, 2)],
+)
+@pytest.mark.asyncio
+async def test_refresh_accounts_owned_session_join_policy(
+    monkeypatch: pytest.MonkeyPatch,
+    join_existing: bool,
+    expected_calls: int,
+) -> None:
+    account = _make_account("acc_owned_join_policy", "workspace_owned_join_policy")
+    stored_account = _make_account("acc_owned_join_policy", "workspace_owned_join_policy")
+    stored_account.status = AccountStatus.PAUSED
+    started = asyncio.Event()
+    release = asyncio.Event()
+    refresh_calls = 0
+
+    @dataclass(frozen=True, slots=True)
+    class Settings:
+        usage_refresh_enabled: bool = True
+        usage_refresh_interval_seconds: int = 0
+        usage_refresh_auth_failure_cooldown_seconds: int = 0
+
+    class AccountsRepo:
+        async def get_by_id(self, account_id: str):
+            return account if account_id == account.id else None
+
+        async def get_by_id_fresh(self, account_id: str):
+            return stored_account if account_id == account.id else None
+
+    async def fake_owned_refresh(
+        self: UsageUpdater,
+        account_id: str,
+        *,
+        interval_seconds: int,
+    ) -> usage_updater_module.AccountRefreshResult:
+        nonlocal refresh_calls
+        assert self is not None
+        assert account_id == account.id
+        assert interval_seconds == 0
+        refresh_calls += 1
+        started.set()
+        await release.wait()
+        return usage_updater_module.AccountRefreshResult(usage_written=False)
+
+    monkeypatch.setattr(usage_updater_module, "get_settings", Settings)
+    monkeypatch.setattr(UsageUpdater, "_refresh_account_if_stale_with_owned_session", fake_owned_refresh)
+
+    updater = UsageUpdater(StubUsageRepository(), AccountsRepo())
+    first = asyncio.create_task(
+        updater.refresh_accounts(
+            [account],
+            {},
+            own_singleflight_sessions=True,
+            join_existing=join_existing,
+        )
+    )
+    await asyncio.wait_for(started.wait(), timeout=1)
+    second = asyncio.create_task(
+        updater.refresh_accounts(
+            [account],
+            {},
+            own_singleflight_sessions=True,
+            join_existing=join_existing,
+        )
+    )
+
+    release.set()
+    await asyncio.gather(first, second)
+
+    assert refresh_calls == expected_calls
+    if join_existing:
+        assert account.status == AccountStatus.PAUSED
+
 @pytest.mark.asyncio
 async def test_owned_singleflight_reload_skips_account_that_became_ineligible(
     monkeypatch: pytest.MonkeyPatch,
@@ -267,6 +348,9 @@ def __init__(self, session) -> None:
         async def get_by_id(self, account_id: str):
             return account if account_id == account.id else None
 
+        async def get_by_id_fresh(self, account_id: str):
+            return account if account_id == account.id else None
+
     @asynccontextmanager
     async def background_session():
         yield object()