diff --git a/app/modules/proxy/_service/rate_limit.py b/app/modules/proxy/_service/rate_limit.py index 3c0619a01f..8f2f988261 100644 --- a/app/modules/proxy/_service/rate_limit.py +++ b/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, diff --git a/app/modules/usage/updater.py b/app/modules/usage/updater.py index 2327ecba14..f375e930e7 100644 --- a/app/modules/usage/updater.py +++ b/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 diff --git a/tests/unit/test_proxy_rate_limit.py b/tests/unit/test_proxy_rate_limit.py index a740eaf344..03a12cb559 100644 --- a/tests/unit/test_proxy_rate_limit.py +++ b/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 diff --git a/tests/unit/test_usage_updater.py b/tests/unit/test_usage_updater.py index bf624e747e..0143373628 100644 --- a/tests/unit/test_usage_updater.py +++ b/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()