diff --git a/.all-contributorsrc b/.all-contributorsrc index 78c8c9726..21dc26a0d 100644 --- a/.all-contributorsrc +++ b/.all-contributorsrc @@ -766,6 +766,16 @@ "code", "test" ] + }, + { + "login": "ellentane", + "name": "Jonáš Sivek", + "avatar_url": "https://avatars.githubusercontent.com/u/70338266?v=4", + "profile": "https://github.com/ellentane", + "contributions": [ + "code", + "test" + ] } ], "contributorsPerLine": 7, diff --git a/README.md b/README.md index 4d5c0e1d7..a816c9df4 100644 --- a/README.md +++ b/README.md @@ -585,6 +585,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/e n3crosis
n3crosis
💻 ⚠️ copilot
copilot
💻 geoHeil
geoHeil
💻 ⚠️ + Jonáš Sivek
Jonáš Sivek
💻 ⚠️ diff --git a/app/core/clients/headers.py b/app/core/clients/headers.py new file mode 100644 index 000000000..bc2758978 --- /dev/null +++ b/app/core/clients/headers.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from app.core.utils.request_id import get_request_id + + +def build_chatgpt_auth_headers( + access_token: str, + account_id: str | None, + *, + extra: dict[str, str] | None = None, +) -> dict[str, str]: + """Build the headers required to call ChatGPT ``backend-api`` endpoints. + + Includes the OAuth bearer token and the ``chatgpt-account-id`` header. The + account-id header is omitted when the id is a synthetic ``email_``/``local_`` + prefix, matching upstream behavior. An active request id (if any) is attached. + """ + headers: dict[str, str] = { + "Authorization": f"Bearer {access_token}", + "Accept": "application/json", + } + request_id = get_request_id() + if request_id: + headers["x-request-id"] = request_id + if account_id and not account_id.startswith(("email_", "local_")): + headers["chatgpt-account-id"] = account_id + if extra: + headers.update(extra) + return headers diff --git a/app/core/clients/rate_limit_reset_credits.py b/app/core/clients/rate_limit_reset_credits.py new file mode 100644 index 000000000..e0eecfd11 --- /dev/null +++ b/app/core/clients/rate_limit_reset_credits.py @@ -0,0 +1,438 @@ +from __future__ import annotations + +import asyncio +import logging +import uuid +from datetime import datetime + +import aiohttp +from aiohttp_retry import ExponentialRetry, RetryClient +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from app.core.clients.codex import ( + CodexClient, + CodexTransportError, + create_codex_session, + require_route_or_direct_egress_opt_in, +) +from app.core.clients.headers import build_chatgpt_auth_headers +from app.core.clients.http import lease_retry_client +from app.core.clients.usage import ( + _retry_delay_seconds, + _safe_codex_json, +) +from app.core.config.settings import get_settings +from app.core.types import JsonObject +from app.core.upstream_proxy import ResolvedUpstreamRoute +from app.core.utils.request_id import get_request_id + +RETRYABLE_STATUS = {408, 429, 500, 502, 503, 504} +RETRY_START_TIMEOUT = 0.5 +RETRY_MAX_TIMEOUT = 2.0 + +logger = logging.getLogger(__name__) + + +class ResetCreditFetchError(Exception): + def __init__(self, status_code: int, message: str, code: str | None = None) -> None: + super().__init__(message) + self.status_code = status_code + self.message = message + self.code = code + + +class ConsumeResetCreditError(Exception): + def __init__(self, status_code: int, message: str, code: str | None = None) -> None: + super().__init__(message) + self.status_code = status_code + self.message = message + self.code = code + + +class ResetCreditItem(BaseModel): + model_config = ConfigDict(extra="ignore") + + id: str + reset_type: str | None = None + status: str | None = None + granted_at: datetime | None = None + expires_at: datetime | None = None + title: str | None = None + description: str | None = None + redeem_started_at: datetime | None = None + redeemed_at: datetime | None = None + + +class ResetCreditsResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + credits: list[ResetCreditItem] + available_count: int + + +class ConsumeResetCreditCredit(BaseModel): + model_config = ConfigDict(extra="ignore") + + id: str | None = None + reset_type: str | None = None + status: str | None = None + redeemed_at: datetime | None = None + + +class ConsumeResetCreditResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + code: str + credit: ConsumeResetCreditCredit + windows_reset: int + + +class RateLimitResetCreditsSnapshot(BaseModel): + """In-memory snapshot of a single account's banked reset credits. + + Carries the upstream ``available_count``, the soonest expiry among the + available credits, and the full credit list so dashboard endpoints can + render details and the consume path can re-select at click time. + """ + + model_config = ConfigDict(extra="ignore") + + available_count: int = 0 + nearest_expires_at: datetime | None = None + credits: list[ResetCreditItem] = Field(default_factory=list) + + +async def fetch_reset_credits( + access_token: str, + account_id: str | None, + *, + base_url: str | None = None, + timeout_seconds: float | None = None, + max_retries: int | None = None, + client: RetryClient | None = None, + route: ResolvedUpstreamRoute | None = None, + codex_client: CodexClient | None = None, + allow_direct_egress: bool = False, +) -> ResetCreditsResponse: + settings = get_settings() + usage_base = base_url or settings.upstream_base_url + url = _reset_credits_url(usage_base) + timeout = aiohttp.ClientTimeout(total=timeout_seconds or settings.usage_fetch_timeout_seconds) + retries = max_retries if max_retries is not None else settings.usage_fetch_max_retries + headers = build_chatgpt_auth_headers(access_token, account_id) + retry_options = _retry_options(retries + 1) + require_route_or_direct_egress_opt_in( + route=route, + allow_direct_egress=allow_direct_egress, + operation="reset credits fetch", + ) + + try: + if route is not None: + data = await _fetch_reset_credits_via_codex( + url=url, + route=route, + headers=headers, + timeout_seconds=timeout_seconds or settings.usage_fetch_timeout_seconds, + retries=retries, + codex_client=codex_client, + ) + else: + async with lease_retry_client(client) as retry_client: + async with retry_client.request( + "GET", + url, + headers=headers, + timeout=timeout, + retry_options=retry_options, + ) as resp: + data = await _safe_json(resp) + if resp.status >= 400: + code = _extract_error_code(data) + message = _extract_error_message(data) or f"Reset credits fetch failed ({resp.status})" + logger.warning( + "Reset credits fetch failed request_id=%s status=%s code=%s message=%s", + get_request_id(), + resp.status, + code, + message, + ) + raise ResetCreditFetchError(resp.status, message, code=code) + try: + return ResetCreditsResponse.model_validate(_success_payload(data)) + except (ValueError, ValidationError) as exc: + logger.warning("Reset credits fetch invalid payload request_id=%s", get_request_id()) + raise ResetCreditFetchError(502, "Invalid reset credits payload") from exc + except (aiohttp.ClientError, asyncio.TimeoutError, CodexTransportError) as exc: + logger.warning("Reset credits fetch error request_id=%s error=%s", get_request_id(), exc) + raise ResetCreditFetchError(0, f"Reset credits fetch failed: {exc}") from exc + + +async def consume_reset_credit( + access_token: str, + account_id: str | None, + credit_id: str, + *, + base_url: str | None = None, + timeout_seconds: float | None = None, + max_retries: int | None = None, + client: RetryClient | None = None, + route: ResolvedUpstreamRoute | None = None, + codex_client: CodexClient | None = None, + allow_direct_egress: bool = False, +) -> ConsumeResetCreditResponse: + settings = get_settings() + usage_base = base_url or settings.upstream_base_url + url = _consume_url(usage_base) + timeout = aiohttp.ClientTimeout(total=timeout_seconds or settings.usage_fetch_timeout_seconds) + # Consume is non-idempotent, so omitted max_retries must not inherit the + # fetch retry budget and risk replaying a successful upstream redemption. + retries = max_retries if max_retries is not None else 0 + headers = build_chatgpt_auth_headers( + access_token, + account_id, + extra={"Content-Type": "application/json"}, + ) + redeem_request_id = str(uuid.uuid4()) + body = {"credit_id": credit_id, "redeem_request_id": redeem_request_id} + retry_options = _retry_options(retries + 1) + require_route_or_direct_egress_opt_in( + route=route, + allow_direct_egress=allow_direct_egress, + operation="reset credits consume", + ) + + try: + if route is not None: + return await _consume_reset_credit_via_codex( + url=url, + route=route, + headers=headers, + body=body, + timeout_seconds=timeout_seconds or settings.usage_fetch_timeout_seconds, + retries=retries, + codex_client=codex_client, + ) + async with lease_retry_client(client) as retry_client: + async with retry_client.request( + "POST", + url, + headers=headers, + json=body, + timeout=timeout, + retry_options=retry_options, + ) as resp: + data = await _safe_json(resp) + if resp.status >= 400: + code = _extract_error_code(data) + message = _extract_error_message(data) or f"Reset credits consume failed ({resp.status})" + logger.warning( + "Reset credits consume failed request_id=%s status=%s code=%s message=%s", + get_request_id(), + resp.status, + code, + message, + ) + raise ConsumeResetCreditError(resp.status, message, code=code) + try: + return ConsumeResetCreditResponse.model_validate(_success_payload(data)) + except (ValueError, ValidationError) as exc: + logger.warning("Reset credits consume invalid payload request_id=%s", get_request_id()) + raise ConsumeResetCreditError(502, "Invalid reset credits consume payload") from exc + except (aiohttp.ClientError, asyncio.TimeoutError, CodexTransportError) as exc: + logger.warning("Reset credits consume error request_id=%s error=%s", get_request_id(), exc) + raise ConsumeResetCreditError(0, f"Reset credits consume failed: {exc}") from exc + + +def build_snapshot(response: ResetCreditsResponse) -> RateLimitResetCreditsSnapshot: + """Project an upstream list response into the cached snapshot shape.""" + nearest = _nearest_available_expires_at(response.credits) + return RateLimitResetCreditsSnapshot( + available_count=response.available_count, + nearest_expires_at=nearest, + credits=list(response.credits), + ) + + +def _nearest_available_expires_at(credits: list[ResetCreditItem]) -> datetime | None: + candidates = [ + credit.expires_at for credit in credits if credit.status == "available" and credit.expires_at is not None + ] + return min(candidates) if candidates else None + + +def _reset_credits_url(base_url: str) -> str: + normalized = base_url.rstrip("/") + if "/backend-api" not in normalized: + normalized = f"{normalized}/backend-api" + return f"{normalized}/wham/rate-limit-reset-credits" + + +def _consume_url(base_url: str) -> str: + return f"{_reset_credits_url(base_url)}/consume" + + +async def _safe_json(resp: aiohttp.ClientResponse) -> JsonObject: + try: + data = await resp.json(content_type=None) + except Exception: + text = await resp.text() + return {"error": {"message": text.strip()}} + return data if isinstance(data, dict) else {"error": {"message": str(data)}} + + +def _success_payload(payload: JsonObject) -> JsonObject: + if "error" in payload: + raise ValueError("success response carried error payload") + return payload + + +class _ErrorEnvelope(BaseModel): + model_config = ConfigDict(extra="ignore") + + error: dict[str, str | None] | str | None = None + error_description: str | None = None + message: str | None = None + + +def _extract_error_message(payload: JsonObject) -> str | None: + envelope = _ErrorEnvelope.model_validate(payload) + error = envelope.error + if isinstance(error, dict): + message = error.get("message") + if isinstance(message, str) and message: + return message + description = error.get("error_description") + if isinstance(description, str) and description: + return description + if isinstance(error, str) and error: + return envelope.error_description or error + return envelope.message + + +def _extract_error_code(payload: JsonObject) -> str | None: + envelope = _ErrorEnvelope.model_validate(payload) + error = envelope.error + if isinstance(error, dict): + code = error.get("code") + if isinstance(code, str): + normalized = code.strip().lower() + return normalized or None + return None + + +def _retry_options(attempts: int) -> ExponentialRetry: + return ExponentialRetry( + attempts=attempts, + start_timeout=RETRY_START_TIMEOUT, + max_timeout=RETRY_MAX_TIMEOUT, + factor=2.0, + statuses=RETRYABLE_STATUS, + exceptions={aiohttp.ClientError, asyncio.TimeoutError}, + retry_all_server_errors=False, + ) + + +async def _fetch_reset_credits_via_codex( + *, + url: str, + route: ResolvedUpstreamRoute, + headers: dict[str, str], + timeout_seconds: float, + retries: int, + codex_client: CodexClient | None, +) -> JsonObject: + attempts = max(1, retries + 1) + owns_codex_client = codex_client is None + active_codex_client = codex_client or CodexClient(create_codex_session()) + try: + for attempt in range(attempts): + try: + resp = await active_codex_client.request( + "GET", + url, + route=route, + headers=headers, + timeout=timeout_seconds, + ) + except CodexTransportError: + if attempt < attempts - 1: + await asyncio.sleep(_retry_delay_seconds(attempt)) + continue + raise + + data = await _safe_codex_json(resp) + status = _codex_response_status(resp) + if status in RETRYABLE_STATUS and attempt < attempts - 1: + await asyncio.sleep(_retry_delay_seconds(attempt)) + continue + if status >= 400: + code = _extract_error_code(data) + message = _extract_error_message(data) or f"Reset credits fetch failed ({status})" + raise ResetCreditFetchError(status, message, code=code) + return data if isinstance(data, dict) else {"error": {"message": str(data)}} + finally: + if owns_codex_client: + close = getattr(active_codex_client, "close", None) + if callable(close): + await close() + raise RuntimeError("unreachable reset credits fetch retry state") + + +async def _consume_reset_credit_via_codex( + *, + url: str, + route: ResolvedUpstreamRoute, + headers: dict[str, str], + body: dict[str, str], + timeout_seconds: float, + retries: int, + codex_client: CodexClient | None, +) -> ConsumeResetCreditResponse: + attempts = max(1, retries + 1) + owns_codex_client = codex_client is None + active_codex_client = codex_client or CodexClient(create_codex_session()) + try: + for attempt in range(attempts): + try: + resp = await active_codex_client.request( + "POST", + url, + route=route, + headers=headers, + json=body, + timeout=timeout_seconds, + ) + except CodexTransportError: + if attempt < attempts - 1: + await asyncio.sleep(_retry_delay_seconds(attempt)) + continue + raise + + data = await _safe_codex_json(resp) + status = _codex_response_status(resp) + if status in RETRYABLE_STATUS and attempt < attempts - 1: + await asyncio.sleep(_retry_delay_seconds(attempt)) + continue + if status >= 400: + code = _extract_error_code(data) + message = _extract_error_message(data) or f"Reset credits consume failed ({status})" + raise ConsumeResetCreditError(status, message, code=code) + try: + return ConsumeResetCreditResponse.model_validate(_success_payload(data)) + except (ValueError, ValidationError) as exc: + logger.warning("Reset credits consume invalid payload request_id=%s", get_request_id()) + raise ConsumeResetCreditError(502, "Invalid reset credits consume payload") from exc + finally: + if owns_codex_client: + close = getattr(active_codex_client, "close", None) + if callable(close): + await close() + raise RuntimeError("unreachable reset credits consume retry state") + + +def _codex_response_status(response: object) -> int: + value = getattr(response, "status_code", getattr(response, "status", None)) + if value is None: + return 0 + return int(value) diff --git a/app/core/clients/usage.py b/app/core/clients/usage.py index 0a8d7e461..ae8342b4e 100644 --- a/app/core/clients/usage.py +++ b/app/core/clients/usage.py @@ -13,6 +13,7 @@ create_codex_session, require_route_or_direct_egress_opt_in, ) +from app.core.clients.headers import build_chatgpt_auth_headers from app.core.clients.http import lease_retry_client from app.core.config.settings import get_settings from app.core.types import JsonObject @@ -220,13 +221,7 @@ def _usage_url(base_url: str) -> str: def _usage_headers(access_token: str, account_id: str | None) -> dict[str, str]: - headers = {"Authorization": f"Bearer {access_token}", "Accept": "application/json"} - request_id = get_request_id() - if request_id: - headers["x-request-id"] = request_id - if account_id and not account_id.startswith(("email_", "local_")): - headers["chatgpt-account-id"] = account_id - return headers + return build_chatgpt_auth_headers(access_token, account_id) async def _safe_json(resp: aiohttp.ClientResponse) -> JsonObject: diff --git a/app/core/config/settings.py b/app/core/config/settings.py index 302362648..37e861adc 100644 --- a/app/core/config/settings.py +++ b/app/core/config/settings.py @@ -193,6 +193,7 @@ class Settings(BaseSettings): usage_fetch_max_retries: int = 2 usage_refresh_enabled: bool = True usage_refresh_interval_seconds: int = Field(default=60, gt=0) + rate_limit_reset_credits_refresh_interval_seconds: int = Field(default=60, gt=0) openai_cache_affinity_max_age_seconds: int = Field(default=1800, gt=0) warmup_model: str = "gpt-5.4-mini" openai_prompt_cache_key_derivation_enabled: bool = True diff --git a/app/core/exceptions.py b/app/core/exceptions.py index 473f59fd8..edc952f8b 100644 --- a/app/core/exceptions.py +++ b/app/core/exceptions.py @@ -82,3 +82,8 @@ class DashboardRateLimitError(AppError): def __init__(self, message: str, *, retry_after: int, code: str | None = None) -> None: self.retry_after = retry_after super().__init__(message, code=code) + + +class DashboardServiceUnavailableError(AppError): + status_code = 503 + code = "service_unavailable" diff --git a/app/core/handlers/exceptions.py b/app/core/handlers/exceptions.py index 04dc949db..e009f1e36 100644 --- a/app/core/handlers/exceptions.py +++ b/app/core/handlers/exceptions.py @@ -20,6 +20,7 @@ DashboardNotFoundError, DashboardPermissionError, DashboardRateLimitError, + DashboardServiceUnavailableError, DashboardValidationError, ProxyAuthError, ProxyModelNotAllowed, @@ -45,6 +46,7 @@ DashboardBadRequestError, DashboardValidationError, DashboardRateLimitError, + DashboardServiceUnavailableError, ) diff --git a/app/core/usage/models.py b/app/core/usage/models.py index f37c08f1e..1a531e517 100644 --- a/app/core/usage/models.py +++ b/app/core/usage/models.py @@ -27,6 +27,12 @@ class CreditsPayload(BaseModel): balance: str | None = None +class RateLimitResetCreditsSummary(BaseModel): + model_config = ConfigDict(extra="ignore") + + available_count: int | None = None + + class AdditionalRateLimitPayload(BaseModel): model_config = ConfigDict(extra="ignore") @@ -44,4 +50,5 @@ class UsagePayload(BaseModel): seat_type: str | None = None rate_limit: RateLimitPayload | None = None credits: CreditsPayload | None = None + rate_limit_reset_credits: RateLimitResetCreditsSummary | None = None additional_rate_limits: list[AdditionalRateLimitPayload] | None = None diff --git a/app/core/usage/reset_credits_refresh_scheduler.py b/app/core/usage/reset_credits_refresh_scheduler.py new file mode 100644 index 000000000..6a47e9661 --- /dev/null +++ b/app/core/usage/reset_credits_refresh_scheduler.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import asyncio +import contextlib +import logging +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field + +from app.core.clients.rate_limit_reset_credits import ( + ResetCreditFetchError, + ResetCreditsResponse, + build_snapshot, + fetch_reset_credits, +) +from app.core.config.settings import get_settings +from app.core.crypto import TokenEncryptor +from app.core.upstream_proxy import ResolvedUpstreamRoute, UpstreamProxyRouteError +from app.db.models import Account, AccountStatus +from app.db.session import get_background_session +from app.modules.accounts.repository import AccountsRepository +from app.modules.rate_limit_reset_credits.store import ( + RateLimitResetCreditsStore, + get_rate_limit_reset_credits_store, +) +from app.modules.usage.updater import _resolve_upstream_route_for_account + +logger = logging.getLogger(__name__) + +_RESET_CREDITS_SKIP_STATUSES = frozenset( + {AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED} +) + +ResetCreditsFetchFn = Callable[..., Awaitable[ResetCreditsResponse]] +ResolveRouteFn = Callable[[Account], Awaitable[ResolvedUpstreamRoute | None]] + + +@dataclass(slots=True) +class RateLimitResetCreditsRefreshScheduler: + interval_seconds: int + _task: asyncio.Task[None] | None = None + _stop: asyncio.Event = field(default_factory=asyncio.Event) + _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + async def start(self) -> None: + if self._task and not self._task.done(): + return + self._stop.clear() + self._task = asyncio.create_task(self._run_loop()) + + async def stop(self) -> None: + if not self._task: + return + self._stop.set() + self._task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._task + self._task = None + + async def _run_loop(self) -> None: + while not self._stop.is_set(): + await self._refresh_once() + try: + await asyncio.wait_for(self._stop.wait(), timeout=self.interval_seconds) + except asyncio.TimeoutError: + continue + + async def _refresh_once(self) -> None: + async with self._lock: + try: + async with get_background_session() as session: + accounts_repo = AccountsRepository(session) + accounts = await accounts_repo.list_accounts() + await refresh_reset_credits_for_accounts( + accounts=accounts, + encryptor=TokenEncryptor(), + store=get_rate_limit_reset_credits_store(), + fetch_fn=fetch_reset_credits, + resolve_route=_resolve_reset_credits_refresh_route, + ) + except Exception: + logger.exception("Reset credits refresh loop failed") + + +async def refresh_reset_credits_for_accounts( + *, + accounts: list[Account], + encryptor: TokenEncryptor, + store: RateLimitResetCreditsStore, + fetch_fn: ResetCreditsFetchFn = fetch_reset_credits, + resolve_route: ResolveRouteFn | None = None, +) -> None: + """Refresh the cached reset-credits snapshot for each eligible account. + + CRITICAL invariant: this function MUST NOT mutate any account's persisted + status. On upstream error it logs and retains the prior cached snapshot + (i.e. it simply skips overwriting the cache) so account-status derivation + stays owned by usage refresh. One account failing must not abort the loop. + """ + for account in accounts: + if account.status in _RESET_CREDITS_SKIP_STATUSES: + continue + if not account.chatgpt_account_id: + continue + await _refresh_account_reset_credits( + account, + encryptor=encryptor, + store=store, + fetch_fn=fetch_fn, + resolve_route=resolve_route, + ) + + +async def _resolve_reset_credits_refresh_route(account: Account) -> ResolvedUpstreamRoute | None: + return await _resolve_upstream_route_for_account(account, operation="usage_refresh") + + +async def _refresh_account_reset_credits( + account: Account, + *, + encryptor: TokenEncryptor, + store: RateLimitResetCreditsStore, + fetch_fn: ResetCreditsFetchFn, + resolve_route: ResolveRouteFn | None = None, +) -> None: + snapshot_generation = store.generation(account.id) + route: ResolvedUpstreamRoute | None = None + if resolve_route is not None: + try: + route = await resolve_route(account) + except UpstreamProxyRouteError as exc: + logger.warning( + "Reset credits refresh upstream proxy route unavailable account_id=%s reason=%s", + account.id, + exc.reason, + ) + return + try: + access_token = encryptor.decrypt(account.access_token_encrypted) + response = await fetch_fn( + access_token, + account.chatgpt_account_id, + route=route, + allow_direct_egress=route is None, + ) + except ResetCreditFetchError as exc: + logger.warning( + "Reset credits refresh failed account_id=%s error=%s", + account.id, + exc, + ) + return + except Exception as exc: + logger.warning( + "Reset credits refresh failed account_id=%s error=%s", + account.id, + exc, + ) + return + + snapshot = build_snapshot(response) + stored = await store.set_if_generation(account.id, snapshot, snapshot_generation) + if not stored: + logger.info( + "Skipped stale reset credits snapshot account_id=%s", + account.id, + ) + + +def build_rate_limit_reset_credits_scheduler() -> RateLimitResetCreditsRefreshScheduler: + settings = get_settings() + return RateLimitResetCreditsRefreshScheduler( + interval_seconds=settings.rate_limit_reset_credits_refresh_interval_seconds, + ) diff --git a/app/main.py b/app/main.py index 393110e54..4f71d98a3 100644 --- a/app/main.py +++ b/app/main.py @@ -41,6 +41,7 @@ from app.core.resilience.bulkhead import BulkheadMiddleware, get_bulkhead from app.core.resilience.memory_monitor import configure as configure_memory_monitor from app.core.usage.refresh_scheduler import build_usage_refresh_scheduler +from app.core.usage.reset_credits_refresh_scheduler import build_rate_limit_reset_credits_scheduler from app.db.session import SessionLocal, close_db, init_background_db, init_db from app.modules.accounts import api as accounts_api from app.modules.api_keys import api as api_keys_api @@ -63,6 +64,7 @@ ) from app.modules.quota_planner import api as quota_planner_api from app.modules.quota_planner.scheduler import build_quota_planner_scheduler +from app.modules.rate_limit_reset_credits import api as rate_limit_reset_credits_api from app.modules.reports import api as reports_api from app.modules.request_logs import api as request_logs_api from app.modules.runtime import api as runtime_api @@ -151,12 +153,14 @@ async def lifespan(app: FastAPI): sticky_session_cleanup_scheduler = build_sticky_session_cleanup_scheduler() quota_planner_scheduler = build_quota_planner_scheduler() auth_guardian_scheduler = build_auth_guardian_scheduler() + rate_limit_reset_credits_scheduler = build_rate_limit_reset_credits_scheduler() await usage_scheduler.start() await api_key_limit_reset_scheduler.start() await model_scheduler.start() await sticky_session_cleanup_scheduler.start() await quota_planner_scheduler.start() await auth_guardian_scheduler.start() + await rate_limit_reset_credits_scheduler.start() if settings.metrics_enabled and PROMETHEUS_AVAILABLE: import uvicorn @@ -317,6 +321,7 @@ async def _activate_bridge_membership(svc: RingMembershipService, iid: str) -> N await model_scheduler.stop() await api_key_limit_reset_scheduler.stop() await usage_scheduler.stop() + await rate_limit_reset_credits_scheduler.stop() try: await close_http_client() finally: @@ -387,6 +392,7 @@ def create_app() -> FastAPI: app.include_router(proxy_api.usage_router) app.include_router(audit_api.router) app.include_router(accounts_api.router) + app.include_router(rate_limit_reset_credits_api.router) app.include_router(dashboard_api.router) app.include_router(usage_api.router) app.include_router(request_logs_api.router) diff --git a/app/modules/accounts/mappers.py b/app/modules/accounts/mappers.py index cff607e40..a99339f44 100644 --- a/app/modules/accounts/mappers.py +++ b/app/modules/accounts/mappers.py @@ -22,9 +22,17 @@ AccountUsageTrend, UsageTrendPoint, ) +from app.modules.rate_limit_reset_credits.store import ( + RateLimitResetCreditsSnapshot, + RateLimitResetCreditsStore, + get_rate_limit_reset_credits_store, +) from app.modules.usage.mappers import usage_history_to_window_row _ACCOUNT_ROUTING_POLICIES = frozenset({"burn_first", "normal", "preserve"}) +_RESET_CREDITS_INELIGIBLE_STATUSES = frozenset( + {AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED} +) _DEFAULT_USAGE_REFRESH_INTERVAL_SECONDS = 60 @@ -39,7 +47,9 @@ def build_account_summaries( limit_warmups_by_account: dict[str, AccountLimitWarmup] | None = None, encryptor: TokenEncryptor, include_auth: bool = True, + reset_credits_store: RateLimitResetCreditsStore | None = None, ) -> list[AccountSummary]: + store = reset_credits_store or get_rate_limit_reset_credits_store() duplicate_keys = _duplicate_detection_keys_appearing_more_than_once(accounts) return [ _account_to_summary( @@ -53,6 +63,7 @@ def build_account_summaries( encryptor, include_auth=include_auth, is_email_duplicate=_duplicate_detection_key(account) in duplicate_keys, + reset_credits_snapshot=_reset_credits_snapshot_for_account(account, store), ) for account in accounts ] @@ -99,6 +110,7 @@ def _account_to_summary( encryptor: TokenEncryptor, include_auth: bool = True, is_email_duplicate: bool = False, + reset_credits_snapshot: RateLimitResetCreditsSnapshot | None = None, ) -> AccountSummary: plan_type = coerce_account_plan_type(account.plan_type, DEFAULT_PLAN) auth_status = _build_auth_status(account, encryptor) if include_auth else None @@ -261,6 +273,8 @@ def _account_to_summary( limit_warmup_enabled=bool(account.limit_warmup_enabled), limit_warmup=_limit_warmup_to_status(limit_warmup), is_email_duplicate=is_email_duplicate, + available_reset_credits=reset_credits_snapshot.available_count if reset_credits_snapshot else 0, + reset_credit_nearest_expires_at=(reset_credits_snapshot.nearest_expires_at if reset_credits_snapshot else None), ) @@ -270,6 +284,15 @@ def _normalize_account_routing_policy(value: str | None) -> str: return "normal" +def _reset_credits_snapshot_for_account( + account: Account, + store: RateLimitResetCreditsStore, +) -> RateLimitResetCreditsSnapshot | None: + if account.status in _RESET_CREDITS_INELIGIBLE_STATUSES or not account.chatgpt_account_id: + return None + return store.get(account.id) + + def _limit_warmup_to_status(entry: AccountLimitWarmup | None) -> AccountLimitWarmupStatus | None: if entry is None: return None diff --git a/app/modules/accounts/schemas.py b/app/modules/accounts/schemas.py index 330032dbb..68bd4c3c5 100644 --- a/app/modules/accounts/schemas.py +++ b/app/modules/accounts/schemas.py @@ -115,6 +115,10 @@ class AccountSummary(DashboardModel): # surface a "delete older" action without requiring the operator to # group rows by email themselves. See codex-lb #787 (B). is_email_duplicate: bool = False + # Banked rate-limit reset credits from the in-memory snapshot when cached, + # otherwise the latest persisted primary usage_history count from /wham/usage. + available_reset_credits: int = 0 + reset_credit_nearest_expires_at: datetime | None = None class AccountsResponse(DashboardModel): diff --git a/app/modules/proxy/api.py b/app/modules/proxy/api.py index eea5e1f0c..27749173a 100644 --- a/app/modules/proxy/api.py +++ b/app/modules/proxy/api.py @@ -5,6 +5,7 @@ import logging import time from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping +from contextlib import asynccontextmanager from datetime import datetime, timezone from json import JSONDecodeError from typing import Any, Final, Literal, cast @@ -37,10 +38,17 @@ validate_proxy_api_key_authorization, validate_usage_api_key, ) +from app.core.auth.refresh import RefreshError from app.core.clients.files import FileProxyError from app.core.clients.proxy import ProxyResponseError +from app.core.clients.rate_limit_reset_credits import ( + ConsumeResetCreditError, + ResetCreditItem, + consume_reset_credit, +) from app.core.config.settings import get_settings from app.core.config.settings_cache import get_settings_cache +from app.core.crypto import TokenEncryptor from app.core.errors import ( PREVIOUS_RESPONSE_STREAM_INCOMPLETE_MESSAGE, OpenAIErrorEnvelope, @@ -78,6 +86,7 @@ from app.core.resilience.overload import is_local_overload_error_code, merge_retry_after_headers from app.core.runtime_logging import log_error_response from app.core.types import JsonValue +from app.core.upstream_proxy import ResolvedUpstreamRoute, UpstreamProxyRouteError, resolve_upstream_route from app.core.utils.json_guards import is_json_mapping from app.core.utils.request_id import get_request_id from app.core.utils.sse import ( @@ -90,6 +99,8 @@ from app.db.models import Account, AccountStatus from app.db.session import get_background_session from app.dependencies import ProxyContext, get_proxy_context, get_proxy_websocket_context +from app.modules.accounts.auth_manager import AuthManager +from app.modules.accounts.repository import AccountsRepository from app.modules.api_keys.repository import ApiKeysRepository from app.modules.api_keys.service import ( TRAFFIC_CLASS_OPPORTUNISTIC, @@ -106,6 +117,7 @@ from app.modules.proxy import affinity as proxy_affinity_module from app.modules.proxy import images_service as images_service_module from app.modules.proxy import service as proxy_service_module +from app.modules.proxy.account_cache import get_account_selection_cache from app.modules.proxy.api_key_usage import estimate_api_key_request_usage from app.modules.proxy.helpers import _rate_limit_details from app.modules.proxy.http_bridge_forwarding import parse_forwarded_request @@ -129,6 +141,9 @@ ModelMetadata, RateLimitStatusPayload, ReasoningLevelSchema, + V1ResetCreditEntry, + V1ResetCreditRedeemRequest, + V1ResetCreditRedeemResponse, V1UsageLimitResponse, V1UsageResponse, WarmupFailedAccount, @@ -142,8 +157,11 @@ RateLimitStatusPayloadData, RateLimitWindowSnapshotData, ) +from app.modules.rate_limit_reset_credits.api import serialize_reset_credit_redeem +from app.modules.rate_limit_reset_credits.store import get_rate_limit_reset_credits_store from app.modules.usage.mappers import usage_history_to_window_row -from app.modules.usage.repository import UsageRepository +from app.modules.usage.repository import AdditionalUsageRepository, UsageRepository +from app.modules.usage.updater import UsageUpdater logger = logging.getLogger(__name__) @@ -168,6 +186,15 @@ ) _PUBLIC_RESPONSES_PRE_CREATED_BUFFER_LIMIT = 64 + +class _V1ResetCreditFreshCredentials: + __slots__ = ("access_token_encrypted", "chatgpt_account_id") + + def __init__(self, *, access_token_encrypted: bytes, chatgpt_account_id: str | None) -> None: + self.access_token_encrypted = access_token_encrypted + self.chatgpt_account_id = chatgpt_account_id + + router = APIRouter( prefix="/backend-api/codex", tags=["proxy"], @@ -711,6 +738,225 @@ async def v1_usage( ) +def _is_reset_credit_selectable_account(account: Account) -> bool: + return bool(account.chatgpt_account_id) and account.status not in ( + AccountStatus.REAUTH_REQUIRED, + AccountStatus.DEACTIVATED, + AccountStatus.PAUSED, + ) + + +def _eligible_reset_credit_accounts(accounts: list[Account], api_key: ApiKeyData) -> list[Account]: + if api_key.account_assignment_scope_enabled: + assigned_ids = {account_id for account_id in api_key.assigned_account_ids if account_id} + requested_accounts = [account for account in accounts if account.id in assigned_ids] + else: + requested_accounts = accounts + return [account for account in requested_accounts if _is_reset_credit_selectable_account(account)] + + +def _project_reset_credit_accounts(accounts: list[Account], api_key: ApiKeyData) -> list[tuple[str, str]]: + eligible_accounts = sorted( + _eligible_reset_credit_accounts(accounts, api_key), + key=lambda account: (account.email, account.id), + ) + return [(account.id, account.email) for account in eligible_accounts] + + +def _list_available_reset_credits(account_id: str, email: str) -> list[V1ResetCreditEntry]: + snapshot = get_rate_limit_reset_credits_store().get(account_id) + if snapshot is None or snapshot.available_count <= 0: + return [] + + available_credits = [credit for credit in snapshot.credits if credit.status == "available"] + if not available_credits: + return [] + + far_future = datetime.max.replace(tzinfo=timezone.utc) + ordered_credits = sorted( + available_credits, + key=lambda credit: (credit.expires_at or far_future, credit.id), + ) + return [ + V1ResetCreditEntry( + account_id=account_id, + email=email, + redeem_id=credit.id, + expired_at=credit.expires_at, + ) + for credit in ordered_credits + ] + + +def _is_reset_credit_account_in_api_key_pool(account: Account | None, api_key: ApiKeyData) -> bool: + if account is None or not _is_reset_credit_selectable_account(account): + return False + if not api_key.account_assignment_scope_enabled: + return True + assigned_ids = {account_id for account_id in api_key.assigned_account_ids if account_id} + return account.id in assigned_ids + + +def _select_available_reset_credit_by_id(account_id: str, redeem_id: str) -> ResetCreditItem | None: + snapshot = get_rate_limit_reset_credits_store().get(account_id) + if snapshot is None or snapshot.available_count <= 0: + return None + for credit in snapshot.credits: + if credit.id == redeem_id and credit.status == "available": + return credit + return None + + +def _translate_v1_reset_credit_consume_error(exc: ConsumeResetCreditError) -> HTTPException: + status_code = exc.status_code if exc.status_code > 0 else 503 + return HTTPException(status_code=status_code, detail=exc.message) + + +def _should_invalidate_v1_reset_credit_snapshot_on_consume_error(exc: ConsumeResetCreditError) -> bool: + return exc.status_code == 409 + + +def _translate_v1_reset_credit_refresh_error(exc: RefreshError) -> HTTPException: + if exc.is_permanent: + get_account_selection_cache().invalidate() + return HTTPException( + status_code=409, + detail=f"Reset credit redeem could not refresh account credentials: {exc.message}", + ) + + +@asynccontextmanager +async def _v1_reset_credit_accounts_refresh_scope() -> AsyncIterator[AccountsRepository]: + async with get_background_session() as session: + yield AccountsRepository(session) + + +async def _ensure_v1_reset_credit_account_fresh(account_id: str) -> _V1ResetCreditFreshCredentials: + async with get_background_session() as session: + repo = AccountsRepository(session) + account = await repo.get_by_id(account_id) + if account is None: + raise HTTPException(status_code=404, detail="Account not found") + auth_manager = AuthManager( + repo, + refresh_repo_factory=_v1_reset_credit_accounts_refresh_scope, + ) + refreshed = await auth_manager.ensure_fresh(account, force=False) + return _V1ResetCreditFreshCredentials( + access_token_encrypted=refreshed.access_token_encrypted, + chatgpt_account_id=refreshed.chatgpt_account_id, + ) + + +@usage_router.get("/v1/reset-credit", response_model=list[V1ResetCreditEntry]) +async def v1_reset_credit( + api_key: ApiKeyData = Security(validate_usage_api_key), +) -> list[V1ResetCreditEntry]: + async with get_background_session() as session: + accounts = await AccountsRepository(session).list_accounts(refresh_existing=True) + eligible_accounts = _project_reset_credit_accounts(accounts, api_key) + + response: list[V1ResetCreditEntry] = [] + for account_id, account_email in eligible_accounts: + response.extend(_list_available_reset_credits(account_id, account_email)) + return response + + +@usage_router.post("/v1/reset-credit", response_model=V1ResetCreditRedeemResponse) +async def v1_redeem_reset_credit( + payload: V1ResetCreditRedeemRequest, + api_key: ApiKeyData = Security(validate_usage_api_key), +) -> V1ResetCreditRedeemResponse: + async with get_background_session() as session: + account = await AccountsRepository(session).get_by_id(payload.account_id) + if not _is_reset_credit_account_in_api_key_pool(account, api_key): + raise HTTPException(status_code=403, detail="Account is outside the API key pool") + if account is None: + raise HTTPException(status_code=403, detail="Account is outside the API key pool") + account_id = account.id + try: + redeem_credentials = await _ensure_v1_reset_credit_account_fresh(account_id) + except RefreshError as exc: + raise _translate_v1_reset_credit_refresh_error(exc) from exc + + async with get_background_session() as session: + account = await AccountsRepository(session).get_by_id(account_id) + if not _is_reset_credit_account_in_api_key_pool(account, api_key): + raise HTTPException(status_code=403, detail="Account is outside the API key pool") + if account is None: + raise HTTPException(status_code=403, detail="Account is outside the API key pool") + try: + route = await _resolve_reset_credit_route(session, account_id) + except UpstreamProxyRouteError as exc: + raise HTTPException(status_code=503, detail="Unable to resolve upstream proxy route") from exc + + async with serialize_reset_credit_redeem(account_id, session=session): + credit = _select_available_reset_credit_by_id(account_id, payload.redeem_id) + if credit is None: + raise HTTPException(status_code=409, detail="Requested reset credit is unavailable") + access_token = TokenEncryptor().decrypt(redeem_credentials.access_token_encrypted) + try: + result = await consume_reset_credit( + access_token, + redeem_credentials.chatgpt_account_id, + credit.id, + route=route, + allow_direct_egress=route is None, + ) + except ConsumeResetCreditError as exc: + if _should_invalidate_v1_reset_credit_snapshot_on_consume_error(exc): + await get_rate_limit_reset_credits_store().invalidate(account_id) + raise _translate_v1_reset_credit_consume_error(exc) from exc + await get_rate_limit_reset_credits_store().invalidate(account_id) + try: + await _refresh_usage_after_v1_reset_credit_redeem(account_id) + except Exception: + logger.warning( + "V1 reset credit consume succeeded but usage refresh failed account_id=%s", + account_id, + exc_info=True, + ) + redeemed_at = result.credit.redeemed_at if result.credit else None + return V1ResetCreditRedeemResponse( + code=result.code, + windows_reset=result.windows_reset, + redeemed_at=redeemed_at, + ) + + +async def _resolve_reset_credit_route(session: AsyncSession, account_id: str) -> ResolvedUpstreamRoute | None: + return await resolve_upstream_route( + session, + account_id=account_id, + operation="reset_credits_consume", + scope="account", + ) + + +async def _refresh_usage_after_v1_reset_credit_redeem(account_id: str) -> None: + async with get_background_session() as session: + account = await AccountsRepository(session).get_by_id(account_id) + if account is None: + logger.warning( + "V1 reset credit consume succeeded but account disappeared before usage refresh account_id=%s", + account_id, + ) + return + usage_updater = UsageUpdater( + UsageRepository(session), + AccountsRepository(session), + AdditionalUsageRepository(session), + ) + refreshed = await usage_updater.force_refresh(account) + if refreshed: + get_account_selection_cache().invalidate() + return + logger.warning( + "V1 reset credit consume succeeded but usage refresh returned no update account_id=%s", + account_id, + ) + + async def _run_v1_warmup( request: Request, context: ProxyContext = Depends(get_proxy_context), diff --git a/app/modules/proxy/schemas.py b/app/modules/proxy/schemas.py index bf48118a0..6d6abaa0b 100644 --- a/app/modules/proxy/schemas.py +++ b/app/modules/proxy/schemas.py @@ -1,5 +1,7 @@ from __future__ import annotations +from datetime import datetime + from pydantic import BaseModel, ConfigDict, Field from app.core.clients.files import OPENAI_FILE_UPLOAD_LIMIT_BYTES, OPENAI_FILE_USE_CASE @@ -227,6 +229,30 @@ class V1UsageResponse(BaseModel): upstream_limits: list[V1UsageLimitResponse] = [] +class V1ResetCreditEntry(BaseModel): + model_config = ConfigDict(extra="forbid") + + account_id: str + email: str + redeem_id: str + expired_at: datetime | None = Field(serialization_alias="expiredAt") + + +class V1ResetCreditRedeemRequest(BaseModel): + model_config = ConfigDict(extra="forbid") + + account_id: str + redeem_id: str + + +class V1ResetCreditRedeemResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + code: str + windows_reset: int + redeemed_at: datetime | None = None + + class WarmupRequest(BaseModel): model_config = ConfigDict(extra="forbid") diff --git a/app/modules/rate_limit_reset_credits/__init__.py b/app/modules/rate_limit_reset_credits/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/app/modules/rate_limit_reset_credits/api.py b/app/modules/rate_limit_reset_credits/api.py new file mode 100644 index 000000000..231f1f1a5 --- /dev/null +++ b/app/modules/rate_limit_reset_credits/api.py @@ -0,0 +1,468 @@ +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Awaitable, Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass +from datetime import datetime, timezone + +from fastapi import APIRouter, Depends, Request +from pydantic import Field +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +from app.core.audit.service import AuditService +from app.core.auth.dependencies import ( + require_dashboard_write_access, + set_dashboard_error_format, + validate_dashboard_session, +) +from app.core.auth.refresh import RefreshError +from app.core.clients.rate_limit_reset_credits import ( + ConsumeResetCreditError, + ConsumeResetCreditResponse, + RateLimitResetCreditsSnapshot, + ResetCreditFetchError, + ResetCreditItem, + ResetCreditsResponse, + build_snapshot, + consume_reset_credit, + fetch_reset_credits, +) +from app.core.crypto import TokenEncryptor +from app.core.exceptions import ( + DashboardAuthError, + DashboardConflictError, + DashboardNotFoundError, + DashboardPermissionError, + DashboardServiceUnavailableError, +) +from app.core.upstream_proxy import ResolvedUpstreamRoute, UpstreamProxyRouteError +from app.core.usage.reset_credits_refresh_scheduler import _refresh_account_reset_credits +from app.db.models import Account, AccountStatus +from app.dependencies import AccountsContext, get_accounts_context +from app.modules.accounts.auth_manager import AuthManager +from app.modules.proxy.account_cache import get_account_selection_cache +from app.modules.rate_limit_reset_credits.store import ( + RateLimitResetCreditsStore, + get_rate_limit_reset_credits_store, +) +from app.modules.shared.schemas import DashboardModel +from app.modules.usage.updater import _resolve_upstream_route_for_account + +logger = logging.getLogger(__name__) + +router = APIRouter( + prefix="/api/accounts", + tags=["dashboard"], + dependencies=[Depends(validate_dashboard_session), Depends(set_dashboard_error_format)], +) + +FetchFn = Callable[..., Awaitable[ResetCreditsResponse]] +ConsumeFn = Callable[..., Awaitable[ConsumeResetCreditResponse]] +RefreshUsageFn = Callable[[Account], Awaitable[None]] +ResolveRouteFn = Callable[[Account], Awaitable[ResolvedUpstreamRoute | None]] + +_NON_REDEEMABLE_STATUSES = frozenset({AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED}) + +_redeem_locks: dict[str, asyncio.Lock] = {} +_redeem_locks_registry_lock = asyncio.Lock() + + +class ResetCreditItemResponse(DashboardModel): + id: str + reset_type: str | None = None + status: str | None = None + granted_at: datetime | None = None + expires_at: datetime | None = None + title: str | None = None + description: str | None = None + redeem_started_at: datetime | None = None + redeemed_at: datetime | None = None + + +class RateLimitResetCreditsSnapshotResponse(DashboardModel): + available_count: int = 0 + nearest_expires_at: datetime | None = None + credits: list[ResetCreditItemResponse] = Field(default_factory=list) + + +class ConsumeResetCreditResponseSchema(DashboardModel): + code: str | None = None + windows_reset: int | None = None + redeemed_at: datetime | None = None + + +@dataclass(frozen=True, slots=True) +class _RedeemResetCreditOutcome: + response: ConsumeResetCreditResponseSchema + available_count_before: int + available_count_after: int + + +@router.get( + "/{account_id}/rate-limit-reset-credits", + response_model=RateLimitResetCreditsSnapshotResponse | None, +) +async def get_rate_limit_reset_credits( + account_id: str, + context: AccountsContext = Depends(get_accounts_context), +) -> RateLimitResetCreditsSnapshotResponse | None: + store = get_rate_limit_reset_credits_store() + account = await context.repository.get_by_id(account_id) + if account is None: + await store.invalidate(account_id) + return None + if account.status in _NON_REDEEMABLE_STATUSES or not account.chatgpt_account_id: + await store.invalidate(account_id) + return None + + snapshot = store.get(account_id) + if snapshot is not None: + return _snapshot_to_response(snapshot) + + await _refresh_account_reset_credits( + account, + encryptor=TokenEncryptor(), + store=store, + fetch_fn=fetch_reset_credits, + resolve_route=_resolve_reset_credit_route, + ) + return _snapshot_to_response(store.get(account_id)) + + +@router.post( + "/{account_id}/rate-limit-reset-credits/consume", + response_model=ConsumeResetCreditResponseSchema, +) +async def consume_rate_limit_reset_credit( + request: Request, + account_id: str, + _write_access=Depends(require_dashboard_write_access), + context: AccountsContext = Depends(get_accounts_context), +) -> ConsumeResetCreditResponseSchema: + account = await context.repository.get_by_id(account_id) + if account is None: + raise DashboardNotFoundError("Account not found", code="account_not_found") + + store = get_rate_limit_reset_credits_store() + + try: + outcome = await _redeem_soonest_reset_credit( + account=account, + store=store, + encryptor=TokenEncryptor(), + lock_session=getattr(context, "session", None), + auth_manager=context.service._auth_manager, + refresh_usage=_build_refresh_usage_callback(context), + resolve_route=_resolve_reset_credit_route, + ) + except RefreshError as exc: + if exc.is_permanent: + get_account_selection_cache().invalidate() + raise DashboardConflictError( + f"Reset credit consume could not refresh account credentials: {exc.message}", + code="account_reset_credit_refresh_failed", + ) from exc + except UpstreamProxyRouteError as exc: + raise DashboardServiceUnavailableError( + f"Reset credit consume upstream proxy route unavailable: {exc.reason}", + code="account_reset_credit_upstream_route_unavailable", + ) from exc + + AuditService.log_async( + "account_rate_limit_reset_credit_consumed", + actor_ip=request.client.host if request.client else None, + details={ + "account_id": account_id, + "consume_code": outcome.response.code, + "windows_reset": outcome.response.windows_reset, + "available_reset_credits_before": outcome.available_count_before, + "available_reset_credits_after": outcome.available_count_after, + }, + ) + return outcome.response + + +async def _redeem_soonest_reset_credit( + *, + account: Account, + store: RateLimitResetCreditsStore, + encryptor: TokenEncryptor, + lock_session: AsyncSession | None = None, + fetch_fn: FetchFn | None = None, + consume_fn: ConsumeFn | None = None, + auth_manager: AuthManager | None = None, + refresh_usage: RefreshUsageFn | None = None, + resolve_route: ResolveRouteFn | None = None, +) -> _RedeemResetCreditOutcome: + _assert_account_can_redeem_reset_credit(account) + effective_fetch_fn = fetch_fn or fetch_reset_credits + effective_consume_fn = consume_fn or consume_reset_credit + + async with serialize_reset_credit_redeem(account.id, session=lock_session): + return await _redeem_soonest_reset_credit_locked( + account=account, + store=store, + encryptor=encryptor, + effective_fetch_fn=effective_fetch_fn, + effective_consume_fn=effective_consume_fn, + auth_manager=auth_manager, + refresh_usage=refresh_usage, + resolve_route=resolve_route, + ) + + +@asynccontextmanager +async def serialize_reset_credit_redeem( + account_id: str, + *, + session: AsyncSession | None, +): + if session is not None and session.get_bind().dialect.name == "postgresql": + await _acquire_postgresql_reset_credit_redeem_lock(session, account_id) + yield + return + + # SQLite and direct unit-test callers keep the existing in-process lock. + lock = await get_reset_credit_redeem_lock(account_id) + async with lock: + yield + + +async def _acquire_postgresql_reset_credit_redeem_lock(session: AsyncSession, account_id: str) -> None: + lock_key = f"reset-credit-redeem:{account_id}" + await session.execute( + text("SELECT pg_advisory_xact_lock(hashtext(:lock_key))"), + {"lock_key": lock_key}, + ) + + +async def _redeem_soonest_reset_credit_locked( + *, + account: Account, + store: RateLimitResetCreditsStore, + encryptor: TokenEncryptor, + effective_fetch_fn: FetchFn, + effective_consume_fn: ConsumeFn, + auth_manager: AuthManager | None, + refresh_usage: RefreshUsageFn | None, + resolve_route: ResolveRouteFn | None, +) -> _RedeemResetCreditOutcome: + redeem_account = account + if auth_manager is not None: + redeem_account = await auth_manager.ensure_fresh(account, force=False) + + cached_snapshot = store.get(account.id) + cached_credit = _select_soonest_available_credit(cached_snapshot) + if cached_credit is None: + raise DashboardConflictError("No available reset credit", code="no_available_reset_credit") + + access_token = encryptor.decrypt(redeem_account.access_token_encrypted) + route: ResolvedUpstreamRoute | None = None + if resolve_route is not None: + route = await resolve_route(redeem_account) + + try: + credits_response = await effective_fetch_fn( + access_token, + redeem_account.chatgpt_account_id, + route=route, + allow_direct_egress=route is None, + ) + except ResetCreditFetchError as exc: + raise _translate_fetch_error(exc) from exc + + credit = _select_soonest_available_credit_from_response(credits_response) + if credit is None: + await store.set(account.id, build_snapshot(credits_response)) + raise DashboardConflictError("No available reset credit", code="no_available_reset_credit") + + try: + result = await effective_consume_fn( + access_token, + redeem_account.chatgpt_account_id, + credit.id, + route=route, + allow_direct_egress=route is None, + ) + except ConsumeResetCreditError as exc: + raise _translate_consume_error(exc) from exc + + redeemed_at = result.credit.redeemed_at if result.credit else None + available_count_after = max(0, credits_response.available_count - 1) + await store.invalidate(account.id) + + if refresh_usage is not None: + try: + await refresh_usage(redeem_account) + except Exception: + logger.warning( + "Reset credit consume succeeded but usage refresh failed account_id=%s", + account.id, + exc_info=True, + ) + await _try_restore_reset_credits_snapshot_after_consume( + account=account, + redeem_account=redeem_account, + encryptor=encryptor, + store=store, + fetch_fn=effective_fetch_fn, + resolve_route=resolve_route, + ) + + return _RedeemResetCreditOutcome( + response=ConsumeResetCreditResponseSchema( + code=result.code, + windows_reset=result.windows_reset, + redeemed_at=redeemed_at, + ), + available_count_before=credits_response.available_count, + available_count_after=available_count_after, + ) + + +def _assert_account_can_redeem_reset_credit(account: Account) -> None: + if account.status in _NON_REDEEMABLE_STATUSES or not account.chatgpt_account_id: + msg = ( + f"Account is {account.status.value} and cannot redeem a reset credit" + if account.status in _NON_REDEEMABLE_STATUSES + else "Account has no ChatGPT account ID and cannot redeem a reset credit" + ) + raise DashboardConflictError( + msg, + code="account_not_reset_credit_applicable", + ) + + +def _build_refresh_usage_callback(context: AccountsContext) -> RefreshUsageFn | None: + usage_updater = context.service._usage_updater + if usage_updater is None: + return None + + async def refresh_usage(account: Account) -> None: + refreshed = await usage_updater.force_refresh(account) + if not refreshed: + raise RuntimeError(f"Forced usage refresh returned no update for account {account.id}") + get_account_selection_cache().invalidate() + + return refresh_usage + + +async def _resolve_reset_credit_route(account: Account) -> ResolvedUpstreamRoute | None: + return await _resolve_upstream_route_for_account(account, operation="rate_limit_reset_consume") + + +async def _try_restore_reset_credits_snapshot_after_consume( + *, + account: Account, + redeem_account: Account, + encryptor: TokenEncryptor, + store: RateLimitResetCreditsStore, + fetch_fn: FetchFn, + resolve_route: ResolveRouteFn | None, +) -> None: + """Best-effort cache repopulation when usage refresh fails after a successful consume.""" + try: + access_token = encryptor.decrypt(redeem_account.access_token_encrypted) + route: ResolvedUpstreamRoute | None = None + if resolve_route is not None: + route = await resolve_route(redeem_account) + credits_response = await fetch_fn( + access_token, + redeem_account.chatgpt_account_id, + route=route, + allow_direct_egress=route is None, + ) + except Exception: + logger.warning( + "Reset credit consume post-refresh re-fetch failed account_id=%s", + account.id, + exc_info=True, + ) + return + await store.set(account.id, build_snapshot(credits_response)) + + +async def get_reset_credit_redeem_lock(account_id: str) -> asyncio.Lock: + lock = _redeem_locks.get(account_id) + if lock is not None: + return lock + async with _redeem_locks_registry_lock: + lock = _redeem_locks.get(account_id) + if lock is None: + lock = asyncio.Lock() + _redeem_locks[account_id] = lock + return lock + + +def _translate_fetch_error(exc: ResetCreditFetchError) -> Exception: + if exc.status_code == 401: + return DashboardAuthError(exc.message, code=exc.code) + if exc.status_code == 403: + return DashboardPermissionError(exc.message, code=exc.code) + if exc.status_code == 409: + return DashboardConflictError(exc.message, code=exc.code) + return DashboardServiceUnavailableError(exc.message, code=exc.code) + + +def _translate_consume_error(exc: ConsumeResetCreditError) -> Exception: + if exc.status_code == 401: + return DashboardAuthError(exc.message, code=exc.code) + if exc.status_code == 403: + return DashboardPermissionError(exc.message, code=exc.code) + if exc.status_code == 409: + return DashboardConflictError(exc.message, code=exc.code) + return DashboardServiceUnavailableError(exc.message, code=exc.code) + + +def _select_soonest_available_credit( + snapshot: RateLimitResetCreditsSnapshot | None, +) -> ResetCreditItem | None: + if snapshot is None: + return None + return _select_soonest_available_credit_from_items(snapshot.credits, snapshot.available_count) + + +def _select_soonest_available_credit_from_response( + response: ResetCreditsResponse, +) -> ResetCreditItem | None: + return _select_soonest_available_credit_from_items(response.credits, response.available_count) + + +def _select_available_credit_by_id( + response: ResetCreditsResponse, + credit_id: str, +) -> ResetCreditItem | None: + if response.available_count <= 0: + return None + for credit in response.credits: + if credit.id == credit_id and credit.status == "available": + return credit + return None + + +def _select_soonest_available_credit_from_items( + credits: list[ResetCreditItem], + available_count: int, +) -> ResetCreditItem | None: + if available_count <= 0: + return None + available = [credit for credit in credits if credit.status == "available"] + if not available: + return None + far_future = datetime.max.replace(tzinfo=timezone.utc) + return min(available, key=lambda credit: credit.expires_at or far_future) + + +def _snapshot_to_response( + snapshot: RateLimitResetCreditsSnapshot | None, +) -> RateLimitResetCreditsSnapshotResponse | None: + if snapshot is None: + return None + return RateLimitResetCreditsSnapshotResponse( + available_count=snapshot.available_count, + nearest_expires_at=snapshot.nearest_expires_at, + credits=[ResetCreditItemResponse.model_validate(credit.model_dump()) for credit in snapshot.credits], + ) diff --git a/app/modules/rate_limit_reset_credits/store.py b/app/modules/rate_limit_reset_credits/store.py new file mode 100644 index 000000000..9fc8aa04d --- /dev/null +++ b/app/modules/rate_limit_reset_credits/store.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from datetime import datetime + +import anyio + +from app.core.clients.rate_limit_reset_credits import RateLimitResetCreditsSnapshot, ResetCreditItem + + +class RateLimitResetCreditsStore: + """In-memory cache of the most recent reset-credits snapshot per account. + + Mirrors the lock-guarded shape of :class:`RateLimitHeadersCache` / + :class:`AccountSelectionCache`. Snapshots are keyed by account id and are + repopulated by each replica's refresh scheduler on every tick; reads from + the dashboard (GET + the AccountSummary mapper) never hit upstream. + """ + + def __init__(self) -> None: + self._snapshots: dict[str, RateLimitResetCreditsSnapshot] = {} + self._lock = anyio.Lock() + self._clear_generation = 0 + self._account_generations: dict[str, int] = {} + + async def set(self, account_id: str, snapshot: RateLimitResetCreditsSnapshot) -> None: + async with self._lock: + self._snapshots[account_id] = snapshot + self._bump_account_generation(account_id) + + def generation(self, account_id: str) -> int: + return self._generation_for(account_id) + + async def set_if_generation( + self, + account_id: str, + snapshot: RateLimitResetCreditsSnapshot, + expected_generation: int, + ) -> bool: + async with self._lock: + if self._generation_for(account_id) != expected_generation: + return False + self._snapshots[account_id] = snapshot + self._bump_account_generation(account_id) + return True + + async def mark_credit_redeemed( + self, + account_id: str, + credit_id: str, + *, + redeemed_at: datetime | None, + ) -> None: + async with self._lock: + snapshot = self._snapshots.get(account_id) + if snapshot is None: + return + updated_credits, matched = _mark_credit_redeemed(snapshot.credits, credit_id, redeemed_at=redeemed_at) + if not matched: + return + available_count = sum(1 for credit in updated_credits if credit.status == "available") + self._snapshots[account_id] = snapshot.model_copy( + update={ + "available_count": available_count, + "nearest_expires_at": _nearest_available_expires_at(updated_credits), + "credits": updated_credits, + } + ) + self._bump_account_generation(account_id) + + def get(self, account_id: str) -> RateLimitResetCreditsSnapshot | None: + return self._snapshots.get(account_id) + + async def invalidate(self, account_id: str | None = None) -> None: + async with self._lock: + if account_id is None: + self._snapshots.clear() + self._clear_generation += 1 + return + self._snapshots.pop(account_id, None) + self._bump_account_generation(account_id) + + def _generation_for(self, account_id: str) -> int: + return self._clear_generation + self._account_generations.get(account_id, 0) + + def _bump_account_generation(self, account_id: str) -> None: + self._account_generations[account_id] = self._account_generations.get(account_id, 0) + 1 + + +_rate_limit_reset_credits_store = RateLimitResetCreditsStore() + + +def get_rate_limit_reset_credits_store() -> RateLimitResetCreditsStore: + return _rate_limit_reset_credits_store + + +def _mark_credit_redeemed( + credits: list[ResetCreditItem], + credit_id: str, + *, + redeemed_at: datetime | None, +) -> tuple[list[ResetCreditItem], bool]: + matched = False + updated: list[ResetCreditItem] = [] + for credit in credits: + if credit.id != credit_id: + updated.append(credit) + continue + matched = True + updated.append(credit.model_copy(update={"status": "redeemed", "redeemed_at": redeemed_at})) + return updated, matched + + +def _nearest_available_expires_at(credits: list[ResetCreditItem]) -> datetime | None: + candidates = [ + credit.expires_at for credit in credits if credit.status == "available" and credit.expires_at is not None + ] + return min(candidates) if candidates else None diff --git a/frontend/src/components/confirm-dialog.test.tsx b/frontend/src/components/confirm-dialog.test.tsx new file mode 100644 index 000000000..26eb525aa --- /dev/null +++ b/frontend/src/components/confirm-dialog.test.tsx @@ -0,0 +1,48 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { ConfirmDialog } from "@/components/confirm-dialog"; + +describe("ConfirmDialog", () => { + it("auto-closes by default after confirm", async () => { + const user = userEvent.setup(); + const onConfirm = vi.fn(); + const onOpenChange = vi.fn(); + + render( + , + ); + + await user.click(screen.getByRole("button", { name: "Confirm" })); + + expect(onConfirm).toHaveBeenCalledTimes(1); + expect(onOpenChange).toHaveBeenCalledWith(false); + }); + + it("stays open when keepOpenOnConfirm is enabled", async () => { + const user = userEvent.setup(); + const onConfirm = vi.fn(); + const onOpenChange = vi.fn(); + + render( + , + ); + + await user.click(screen.getByRole("button", { name: "Confirm" })); + + expect(onConfirm).toHaveBeenCalledTimes(1); + expect(onOpenChange).not.toHaveBeenCalledWith(false); + }); +}); diff --git a/frontend/src/components/confirm-dialog.tsx b/frontend/src/components/confirm-dialog.tsx index 5fbad5852..11fb9323e 100644 --- a/frontend/src/components/confirm-dialog.tsx +++ b/frontend/src/components/confirm-dialog.tsx @@ -14,6 +14,8 @@ export type ConfirmDialogProps = { title: string; description?: string; confirmLabel?: string; + confirmDisabled?: boolean; + keepOpenOnConfirm?: boolean; cancelLabel?: string; onConfirm: () => void; onOpenChange: (open: boolean) => void; @@ -25,6 +27,8 @@ export function ConfirmDialog({ title, description, confirmLabel = "Confirm", + confirmDisabled = false, + keepOpenOnConfirm = false, cancelLabel = "Cancel", onConfirm, onOpenChange, @@ -40,7 +44,17 @@ export function ConfirmDialog({ {children} {cancelLabel} - {confirmLabel} + { + if (keepOpenOnConfirm) { + event.preventDefault(); + } + onConfirm(); + }} + > + {confirmLabel} + diff --git a/frontend/src/components/layout/app-header.test.tsx b/frontend/src/components/layout/app-header.test.tsx new file mode 100644 index 000000000..60d04eab8 --- /dev/null +++ b/frontend/src/components/layout/app-header.test.tsx @@ -0,0 +1,83 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import { HttpResponse, http } from "msw"; +import { MemoryRouter } from "react-router-dom"; +import { describe, expect, it, vi } from "vitest"; + +import { AppHeader } from "@/components/layout/app-header"; +import { server } from "@/test/mocks/server"; +import { createAccountSummary } from "@/test/mocks/factories"; + +function renderHeader() { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + return render( + + + + + , + ); +} + +describe("AppHeader", () => { + it("shows the summed Accounts reset-credit badge capped at 99+", async () => { + server.use( + http.get("/api/accounts", () => + HttpResponse.json({ + accounts: [ + createAccountSummary({ availableResetCredits: 70 }), + createAccountSummary({ accountId: "acc-2", availableResetCredits: 40 }), + ], + }), + ), + ); + + renderHeader(); + + expect(await screen.findAllByText("99+")).not.toHaveLength(0); + }); + + it("sums reset-credit badge across accounts and treats missing counts as zero", async () => { + server.use( + http.get("/api/accounts", () => + HttpResponse.json({ + accounts: [ + createAccountSummary({ availableResetCredits: 5 }), + createAccountSummary({ accountId: "acc-2" }), + createAccountSummary({ accountId: "acc-3", availableResetCredits: null }), + createAccountSummary({ accountId: "acc-4", availableResetCredits: 3 }), + ], + }), + ), + ); + + renderHeader(); + + expect(await screen.findAllByText("8")).not.toHaveLength(0); + }); + + it("hides the Accounts reset-credit badge when no resets are available", async () => { + server.use( + http.get("/api/accounts", () => + HttpResponse.json({ + accounts: [ + createAccountSummary({ availableResetCredits: 0 }), + createAccountSummary({ accountId: "acc-2", availableResetCredits: 0 }), + ], + }), + ), + ); + + renderHeader(); + + await screen.findByRole("link", { name: /Accounts/i }); + expect(screen.queryByText("99+")).not.toBeInTheDocument(); + }); +}); diff --git a/frontend/src/components/layout/app-header.tsx b/frontend/src/components/layout/app-header.tsx index b83a6cd48..7ceada881 100644 --- a/frontend/src/components/layout/app-header.tsx +++ b/frontend/src/components/layout/app-header.tsx @@ -1,3 +1,4 @@ +import { useQuery } from "@tanstack/react-query"; import { Eye, EyeOff, LogIn, LogOut, Menu } from "lucide-react"; import { useState } from "react"; import { NavLink } from "react-router-dom"; @@ -5,6 +6,7 @@ import { NavLink } from "react-router-dom"; import { CodexLogo } from "@/components/brand/codex-logo"; import { Button } from "@/components/ui/button"; import { Sheet, SheetContent, SheetHeader, SheetTitle, SheetTrigger } from "@/components/ui/sheet"; +import { listAccounts } from "@/features/accounts/api"; import { usePrivacyStore } from "@/hooks/use-privacy"; import { cn } from "@/lib/utils"; @@ -35,6 +37,23 @@ export function AppHeader({ const blurred = usePrivacyStore((s) => s.blurred); const togglePrivacy = usePrivacyStore((s) => s.toggle); const PrivacyIcon = blurred ? EyeOff : Eye; + const { data: accounts = [] } = useQuery({ + queryKey: ["accounts", "list"], + queryFn: listAccounts, + select: (data) => data.accounts, + refetchInterval: 30_000, + refetchIntervalInBackground: false, + staleTime: 30_000, + }); + const totalAvailableResetCredits = accounts.reduce( + (total, account) => total + Math.max(0, account.availableResetCredits ?? 0), + 0, + ); + const accountsResetBadge = totalAvailableResetCredits > 99 + ? "99+" + : totalAvailableResetCredits > 0 + ? String(totalAvailableResetCredits) + : null; return (
- {item.label} + + {item.label} + {item.to === "/accounts" && accountsResetBadge ? ( + + {accountsResetBadge} + + ) : null} + ))} @@ -133,13 +159,18 @@ export function AppHeader({ {({ isActive }) => ( {item.label} + {item.to === "/accounts" && accountsResetBadge ? ( + + {accountsResetBadge} + + ) : null} )} diff --git a/frontend/src/features/accounts/api.ts b/frontend/src/features/accounts/api.ts index ba79bde55..9aa5d0b01 100644 --- a/frontend/src/features/accounts/api.ts +++ b/frontend/src/features/accounts/api.ts @@ -15,6 +15,7 @@ import { AccountTrendsResponseSchema, AccountProbeRequestSchema, AccountProbeResponseSchema, + ConsumeRateLimitResetCreditResponseSchema, ManualOauthCallbackRequestSchema, ManualOauthCallbackResponseSchema, OauthCompleteRequestSchema, @@ -22,6 +23,7 @@ import { OauthStartRequestSchema, OauthStartResponseSchema, OauthStatusResponseSchema, + RateLimitResetCreditsSnapshotSchema, RuntimeConnectAddressResponseSchema, } from "@/features/accounts/schemas"; import type { AccountRoutingPolicy } from "@/features/accounts/schemas"; @@ -117,6 +119,20 @@ export function exportAccountAuth(accountId: string) { ); } +export function getRateLimitResetCredits(accountId: string) { + return get( + `${ACCOUNTS_BASE_PATH}/${encodeURIComponent(accountId)}/rate-limit-reset-credits`, + RateLimitResetCreditsSnapshotSchema.nullable(), + ); +} + +export function consumeRateLimitResetCredit(accountId: string) { + return post( + `${ACCOUNTS_BASE_PATH}/${encodeURIComponent(accountId)}/rate-limit-reset-credits/consume`, + ConsumeRateLimitResetCreditResponseSchema, + ); +} + export function deleteAccount(accountId: string, deleteHistory = false) { const qs = deleteHistory ? "?delete_history=true" : ""; return del( diff --git a/frontend/src/features/accounts/components/account-actions.test.tsx b/frontend/src/features/accounts/components/account-actions.test.tsx index edc16b49d..54df9afc1 100644 --- a/frontend/src/features/accounts/components/account-actions.test.tsx +++ b/frontend/src/features/accounts/components/account-actions.test.tsx @@ -20,6 +20,7 @@ describe("AccountActions", () => { onDelete={vi.fn()} onReauth={vi.fn()} onExportAuth={vi.fn()} + onResetCredit={vi.fn()} onSecurityWorkAuthorizedChange={vi.fn()} onLimitWarmupChange={vi.fn()} onRoutingPolicyChange={onRoutingPolicyChange} @@ -46,6 +47,7 @@ describe("AccountActions", () => { onDelete={vi.fn()} onReauth={onReauth} onExportAuth={vi.fn()} + onResetCredit={vi.fn()} onSecurityWorkAuthorizedChange={vi.fn()} onLimitWarmupChange={vi.fn()} onRoutingPolicyChange={vi.fn()} @@ -78,6 +80,7 @@ describe("AccountActions", () => { onDelete={vi.fn()} onReauth={vi.fn()} onExportAuth={vi.fn()} + onResetCredit={vi.fn()} onSecurityWorkAuthorizedChange={vi.fn()} onLimitWarmupChange={vi.fn()} onRoutingPolicyChange={vi.fn()} @@ -107,6 +110,7 @@ describe("AccountActions", () => { onDelete={vi.fn()} onReauth={vi.fn()} onExportAuth={vi.fn()} + onResetCredit={vi.fn()} onSecurityWorkAuthorizedChange={vi.fn()} onLimitWarmupChange={vi.fn()} onRoutingPolicyChange={vi.fn()} @@ -138,6 +142,7 @@ describe("AccountActions", () => { onDelete={vi.fn()} onReauth={vi.fn()} onExportAuth={vi.fn()} + onResetCredit={vi.fn()} onSecurityWorkAuthorizedChange={vi.fn()} onLimitWarmupChange={vi.fn()} onRoutingPolicyChange={vi.fn()} @@ -151,4 +156,95 @@ describe("AccountActions", () => { expect(onProbe).not.toHaveBeenCalled(); }); + + it("shows reset action when reset credits are available", async () => { + const user = userEvent.setup(); + const onResetCredit = vi.fn(); + const account = createAccountSummary({ + availableResetCredits: 3, + resetCreditNearestExpiresAt: "2026-01-03T12:00:00.000Z", + }); + + render( + , + ); + + await user.click(screen.getByRole("button", { name: "Reset (3)" })); + + expect(onResetCredit).toHaveBeenCalledWith(account.accountId); + }); + + it.each(["paused", "deactivated", "reauth_required"] as const)( + "disables reset action for %s accounts", + async (status) => { + const user = userEvent.setup(); + const onResetCredit = vi.fn(); + const account = createAccountSummary({ + status, + availableResetCredits: 2, + resetCreditNearestExpiresAt: "2026-01-03T12:00:00.000Z", + }); + + render( + , + ); + + const button = screen.getByRole("button", { name: "Reset (2)" }); + expect(button).toBeDisabled(); + await user.click(button); + expect(onResetCredit).not.toHaveBeenCalled(); + }, + ); + + it("hides reset action when no reset credits are available", () => { + const account = createAccountSummary({ + availableResetCredits: 0, + resetCreditNearestExpiresAt: null, + }); + + render( + , + ); + + expect(screen.queryByRole("button", { name: /Reset \(/ })).not.toBeInTheDocument(); + }); }); diff --git a/frontend/src/features/accounts/components/account-actions.tsx b/frontend/src/features/accounts/components/account-actions.tsx index 6092b0ea8..a66f73ac0 100644 --- a/frontend/src/features/accounts/components/account-actions.tsx +++ b/frontend/src/features/accounts/components/account-actions.tsx @@ -4,6 +4,7 @@ import { Pause, Play, RefreshCw, + RotateCcw, Route, ShieldCheck, Trash2, @@ -23,6 +24,7 @@ import type { AccountRoutingPolicy, AccountSummary, } from "@/features/accounts/schemas"; +import { formatSingleUnitRemaining } from "@/utils/formatters"; export type AccountActionsProps = { account: AccountSummary; @@ -34,6 +36,7 @@ export type AccountActionsProps = { onDelete: (accountId: string) => void; onReauth: () => void; onExportAuth: (accountId: string) => void; + onResetCredit: (accountId: string) => void; onSecurityWorkAuthorizedChange: (accountId: string, enabled: boolean) => void; onLimitWarmupChange: (accountId: string, enabled: boolean) => void; onRoutingPolicyChange: ( @@ -52,6 +55,7 @@ export function AccountActions({ onDelete, onReauth, onExportAuth, + onResetCredit, onSecurityWorkAuthorizedChange, onLimitWarmupChange, onRoutingPolicyChange, @@ -60,6 +64,16 @@ export function AccountActions({ account.status === "reauth_required" || account.status === "deactivated"; const probeDisabled = busy || readOnly || account.status === "paused" || showOperatorRecoveryAction; + const resetCountdown = account.resetCreditNearestExpiresAt + ? formatSingleUnitRemaining(account.resetCreditNearestExpiresAt) + : null; + const availableResetCredits = account.availableResetCredits ?? 0; + const hasResetCredits = availableResetCredits > 0; + const resetCreditDisabled = + busy || + readOnly || + account.status === "paused" || + showOperatorRecoveryAction; return (
@@ -191,6 +205,33 @@ export function AccountActions({ Export + {hasResetCredits ? ( + + ) : null} + + {hasResetCredits ? ( + + ) : null} {status === "paused" && ( + {hasResetCredits ? ( + + ) : null}
); } diff --git a/frontend/src/utils/formatters.test.ts b/frontend/src/utils/formatters.test.ts index 7258a84c8..6937e4f9c 100644 --- a/frontend/src/utils/formatters.test.ts +++ b/frontend/src/utils/formatters.test.ts @@ -7,6 +7,7 @@ import { formatDateTimeInline, formatAccessTokenLabel, formatCachedTokensMeta, + formatLocalDateTimeSeconds, formatCompactNumber, formatCountdown, formatCurrency, @@ -20,6 +21,7 @@ import { formatQuotaResetMeta, formatRate, formatResetRelative, + formatSingleUnitRemaining, formatRefreshTokenLabel, formatRelative, formatTimeLong, @@ -115,6 +117,15 @@ describe("formatters", () => { expect(formatChartDateTime(iso)).not.toMatch(/AM|PM/); }); + it("formats local timestamps as yyyy-mm-dd hh:mm:ss", () => { + const iso = "2026-01-01T00:00:00.000Z"; + const local = new Date(iso); + const expected = `${local.getFullYear()}-${String(local.getMonth() + 1).padStart(2, "0")}-${String(local.getDate()).padStart(2, "0")} ${String(local.getHours()).padStart(2, "0")}:${String(local.getMinutes()).padStart(2, "0")}:${String(local.getSeconds()).padStart(2, "0")}`; + + expect(formatLocalDateTimeSeconds(iso)).toBe(expected); + expect(formatLocalDateTimeSeconds("bad-date")).toBe("--"); + }); + it("formats relative and countdown values", () => { expect(formatRelative(30 * 60_000)).toBe("in 30m"); expect(formatRelative(90 * 60_000)).toBe("in 2h"); @@ -125,6 +136,29 @@ describe("formatters", () => { expect(formatCountdown(125)).toBe("2:05"); }); + it("formats single-unit reset-credit countdowns", () => { + expect(formatSingleUnitRemaining("2026-01-08T00:00:00.000Z")).toEqual({ + label: "7d", + expiringSoon: false, + }); + expect(formatSingleUnitRemaining("2026-01-07T00:00:00.000Z")).toEqual({ + label: "6d", + expiringSoon: true, + }); + expect(formatSingleUnitRemaining("2026-01-01T01:00:00.000Z")).toEqual({ + label: "1h", + expiringSoon: true, + }); + expect(formatSingleUnitRemaining("2026-01-01T00:01:00.000Z")).toEqual({ + label: "1m", + expiringSoon: true, + }); + expect(formatSingleUnitRemaining("2025-12-31T23:59:59.000Z")).toEqual({ + label: "now", + expiringSoon: true, + }); + }); + it("formats quota reset labels", () => { const in30m = new Date(Date.now() + 30 * 60_000).toISOString(); const in4h13m = new Date(Date.now() + (4 * 60 + 13) * 60_000).toISOString(); diff --git a/frontend/src/utils/formatters.ts b/frontend/src/utils/formatters.ts index 864d3135c..b49623b6b 100644 --- a/frontend/src/utils/formatters.ts +++ b/frontend/src/utils/formatters.ts @@ -226,6 +226,18 @@ export function formatDateTimeInline(iso: string | null | undefined): string { return formatted.time === "--" ? "--" : `${formatted.time} ${formatted.date}`; } +function padTwo(value: number): string { + return String(value).padStart(2, "0"); +} + +export function formatLocalDateTimeSeconds(iso: string | null | undefined): string { + const date = parseDate(iso); + if (!date) { + return "--"; + } + return `${date.getFullYear()}-${padTwo(date.getMonth() + 1)}-${padTwo(date.getDate())} ${padTwo(date.getHours())}:${padTwo(date.getMinutes())}:${padTwo(date.getSeconds())}`; +} + export function formatChartDateTime(iso: string | null | undefined): string { const date = parseDate(iso); return date ? getChartDateTimeFormatter().format(date) : "--"; @@ -285,6 +297,35 @@ export function formatQuotaResetLabel(resetAt: string | null | undefined): strin return formatResetRelative(diffMs); } +const DAY_MS = 86_400_000; +const HOUR_MS = 3_600_000; +const MINUTE_MS = 60_000; +const EXPIRING_SOON_THRESHOLD_MS = 7 * DAY_MS; + +export type SingleUnitRemaining = { + label: string; + expiringSoon: boolean; +}; + +export function formatSingleUnitRemaining(expiresAtIso: string): SingleUnitRemaining { + const ms = new Date(expiresAtIso).getTime() - Date.now(); + if (ms <= 0) { + return { label: "now", expiringSoon: true }; + } + const days = Math.floor(ms / DAY_MS); + const hours = Math.floor(ms / HOUR_MS); + const minutes = Math.floor(ms / MINUTE_MS); + const label = + days >= 1 + ? `${days}d` + : hours >= 1 + ? `${hours}h` + : minutes >= 1 + ? `${minutes}m` + : "now"; + return { label, expiringSoon: ms < EXPIRING_SOON_THRESHOLD_MS }; +} + export function formatQuotaResetMeta( resetAtSecondary: string | null | undefined, windowMinutesSecondary: unknown, diff --git a/openspec/changes/add-rate-limit-reset-credits/.openspec.yaml b/openspec/changes/add-rate-limit-reset-credits/.openspec.yaml new file mode 100644 index 000000000..3ac681e39 --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-06-17 diff --git a/openspec/changes/add-rate-limit-reset-credits/design.md b/openspec/changes/add-rate-limit-reset-credits/design.md new file mode 100644 index 000000000..697a06328 --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/design.md @@ -0,0 +1,69 @@ +## Context + +codex-lb already polls upstream `GET /wham/usage` per account on a 60s leader-gated loop (`app/core/usage/refresh_scheduler.py`) and uses an in-memory cache pattern (`app/modules/proxy/rate_limit_cache.py`) for short-lived values. Per-account OAuth bearer tokens are stored encrypted at rest (`Account.access_token_encrypted`) and decrypted on demand via `TokenEncryptor`. Dashboard endpoints authenticate via `validate_dashboard_session` + `require_dashboard_write_access`. + +OpenAI added banked rate-limit reset credits on 2026-06-12. The upstream API surface (reverse-engineered and documented in the [`aaamosh/codex-reset`](https://github.com/aaamosh/codex-reset) reference) is: + +- `GET /wham/rate-limit-reset-credits` → `{credits: [...], available_count: N}` with `Authorization: Bearer ` + `chatgpt-account-id: ` headers +- `POST /wham/rate-limit-reset-credits/consume` with body `{credit_id, redeem_request_id}` → `{code, credit, windows_reset}` + +The reference is a single-account CLI; codex-lb needs a multi-account, dashboard-driven, in-memory-cached variant. + +## Goals / Non-Goals + +**Goals:** +- Per-account background poll of reset credits every 60s, cached in-memory keyed by account id +- Dashboard operators can see, per account: how many banked credits are available and when the soonest one expires +- Dashboard operators can redeem the soonest-expiring credit for any account from three surfaces (Accounts action bar, Dashboard table, Dashboard grid) with a confirmation dialog +- Dashboard operators can see the summed available reset-credit count on the top-nav Accounts tab +- Sort the Accounts page by available reset credits +- Reuse existing token decryption, scheduler shape, in-memory cache shape, dashboard auth, confirmation dialog, and formatter conventions — no new frameworks + +**Non-Goals:** +- No DB persistence — snapshots live only in memory and are repopulated after restart +- No referral/invite logic from the reference repo +- No changes to `/wham/usage` or account status derivation (rate_limited / quota_exceeded reconciliation stays owned by usage refresh) +- No live-ticking countdown — values recompute on each 60s scheduler tick + TanStack Query refetch, matching the existing `formatQuotaResetLabel` pattern +- No new top-nav account card; the badge lives on `AccountListItem` only +- No mobile-specific behavior + +## Decisions + +### Decision: Dedicated module + scheduler, mirroring `usage_refresh` (not folding into the usage loop) +**Rationale:** Usage refresh owns account-status derivation and has a dense, scenario-heavy spec (`usage-refresh-policy`) tying it to quota reconciliation, cooldowns, and warm-up. Bolting credits onto that loop would couple two upstream calls and their failure modes, and muddy the usage-refresh contract. A dedicated `RateLimitResetCreditsRefreshScheduler` reuses the `UsageRefreshScheduler` loop shape (`asyncio.Lock`-guarded `_refresh_once`, interval-only configuration) and always starts with the application. Unlike usage refresh, it deliberately runs on every replica because reset-credit snapshots are process-local and dashboard reads must be consistent regardless of which replica handles the request. +**Alternatives considered:** (a) Fold into `UsageRefreshScheduler._refresh_once` — rejected for the coupling above. (b) Pure passthrough via the local `wham_router` proxy — rejected because the dashboard needs the in-memory store and per-account token decryption that the proxy router does not have, and the requirement is "refresh every 60s + store in-memory." + +### Decision: Server picks the soonest-expiring credit at consume time +**Rationale:** Single source of truth. The client passes only `{account_id}` to `POST /consume`; the server reads the cached snapshot, selects the available credit with the smallest `expires_at`, generates `redeem_request_id = uuid4()`, and calls upstream. This guarantees "nearest expiry_at is selected" even if the UI is stale, and avoids a client/server clock skew race. +**Alternatives considered:** Client sends the specific `credit_id` — rejected because the cached snapshot may have changed between render and click (e.g. one expired or was redeemed elsewhere). + +### Decision: Expose `available_reset_credits` + `reset_credit_nearest_expires_at` on `AccountSummary` (no DB column) +**Rationale:** The Accounts-page and Dashboard list both consume `AccountSummary`; joining the cached snapshot at mapper time gets the data to every UI surface with one change and zero migration. Account rows that have no cache yet return `0` / `null` and the UI hides its reset affordances for them. +**Alternatives considered:** Separate `/api/accounts/{id}/rate-limit-reset-credits` GET consumed per-card — rejected because it adds N round-trips and N re-renders; the count belongs on the summary the UI already fetches. + +### Decision: Countdown is single-unit and goes red under 7 days +**Rationale:** User requirement: one unit only ("6d" / "13h" / "45m" / "now"), red when `< 7d`. A new `formatSingleUnitRemaining(expiresAtIso)` helper sits next to existing `formatQuotaResetLabel` / `formatResetRelative` in `utils/formatters.ts`; the caller colors it via `ms < 7 * DAY_MS`. We do NOT add a ticking hook (matches the existing reset-label pattern). The confirmation dialog uses a separate local-time formatter with exact `YYYY-MM-DD HH:MM:SS` output so operators can see the precise expiry instant in their own timezone. +**Alternatives considered:** Reuse `formatResetRelative` — rejected because it returns multi-unit ("6d 13h") output. + +### Decision: Reset credit refresh never mutates account status +**Rationale:** Account status (active / rate_limited / quota_exceeded / paused / deactivated) is owned by usage refresh. Reset-credit polling failure MUST NOT deactivate or block an account — doing so would create a second status owner and contradict `usage-refresh-policy`. On upstream errors the scheduler logs, keeps the prior snapshot if any, and moves on. +**Alternatives considered:** Reuse usage-refresh cooldown/deactivation classification — rejected because it would require this scheduler to write account status, violating the single-owner invariant. + +## Risks / Trade-offs + +- **[Upstream endpoints are undocumented]** → Mitigation: client treats non-200 / non-JSON defensively, logs, keeps prior snapshot; consume-failure surfaces to UI as a toast without invalidating the cache. Document the upstream-dependence caveat in the capability `context.md`. +- **[In-memory cache lost on restart]** → Mitigation: acceptable per requirements; the next 60s tick repopulates. UI treats missing snapshot as `available_reset_credits: 0` (hidden affordances), not an error. +- **[Credit consumed even on partial reset]** (upstream behavior: "if POST returns 200, the credit is gone") → Mitigation: require an explicit confirmation step before redeeming. On success we invalidate the cache and let the next tick reconcile. +- **[Race: credit expires between render and click]** → Mitigation: server re-selects from the freshest cached snapshot at consume time and surfaces upstream's error if the chosen credit is no longer redeemable. +- **[Many accounts = many upstream calls per tick]** → Mitigation: reuse the same skip rules (paused/deactivated/missing chatgpt-account-id) and keep the interval configurable. Each replica polls so its process-local cache is useful for dashboard reads; moving snapshots to shared storage can later reduce duplicate polling if upstream load becomes a problem. +- **[Guest read-only dashboard users]** → Mitigation: `POST /consume` requires `require_dashboard_write_access`; guests can see the badge/button (count is read off `AccountSummary`) but cannot redeem. + +## Migration Plan + +- No DB migration (in-memory only). No env var required beyond the refresh interval setting. +- Deploy is a single rolling restart; the first tick after boot repopulates snapshots within 60s. +- Rollback: revert the deploy; there is no separate disable toggle for reset-credit polling. + +## Open Questions + +None blocking. (The two upstream endpoints' longevity is a runtime risk documented in `context.md`, not a design unknown.) diff --git a/openspec/changes/add-rate-limit-reset-credits/proposal.md b/openspec/changes/add-rate-limit-reset-credits/proposal.md new file mode 100644 index 000000000..d85f70a35 --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/proposal.md @@ -0,0 +1,32 @@ +## Why + +OpenAI rolled out savable ("banked") rate-limit reset credits for Codex on 2026-06-12. Eligible ChatGPT plans receive credits that can be redeemed to reset rate-limit windows, but the redeem affordance only ships in the desktop app and VS Code/Cursor/Windsurf extension — not in the Codex CLI, and not in any operator surface. codex-lb operators managing many accounts have no way to see how many banked resets an account has, when they expire, or to redeem one without leaving the dashboard. We need first-class visibility and a one-click redeem action that reuses each account's existing OAuth bearer token. + +## What Changes + +- Add a per-account background poller that calls upstream `GET /wham/rate-limit-reset-credits` every 60s using each account's stored bearer token, and caches the result in-memory keyed by account id. +- Add a dashboard endpoint `POST /api/accounts/{account_id}/rate-limit-reset-credits/consume` that redeems the soonest-expiring available credit by calling upstream `POST /wham/rate-limit-reset-credits/consume` with `{credit_id, redeem_request_id}`. +- Expose `available_reset_credits` count and `reset_credit_nearest_expires_at` timestamp on the account summary payloads consumed by the Accounts page and Dashboard. +- Accounts page: add a `Reset (N)` button to the per-account action bar (next to Export), a count badge on each `AccountListItem` (capped at "99+"), and a new "Most reset credits" option in the sort-mode dropdown. +- Dashboard Accounts section: add a reset action next to Details in both the table and grid views, with the grid label rendered as `Reset (N)`. +- Show a single-unit countdown ("6d" / "13h" / "45m" / "now") of the nearest credit's expiry on each Reset button; render it in destructive/red when less than 7 days remain. +- Show the confirmation-dialog expiry in local time using `YYYY-MM-DD HH:MM:SS`. +- Show the total available reset-credit count on the top-nav Accounts tab, pinned to the upper-right radius and capped at `99+`. +- Gate the Reset button and badge on `available_reset_credits > 0`; gate the redeem endpoint on dashboard write access (read-only guests cannot redeem). + +## Capabilities + +### New Capabilities +- `rate-limit-reset-credits`: Background polling, in-memory caching, and dashboard-initiated redemption of upstream Codex banked rate-limit reset credits per account. + +### Modified Capabilities +- `frontend-architecture`: New dashboard/Accounts UI elements — Reset button (Accounts tab + Dashboard), count badge on AccountListItem, top-nav Accounts total badge, expiry countdown label, local expiry timestamp formatting, and a new Accounts sort mode by available reset credits. + +## Impact + +- **Backend (new)**: `app/core/clients/rate_limit_reset_credits.py` (upstream client), `app/modules/rate_limit_reset_credits/store.py` (in-memory store + singleton), `app/core/usage/reset_credits_refresh_scheduler.py` (60s per-replica loop), `app/modules/rate_limit_reset_credits/api.py` (dashboard GET + consume POST). +- **Backend (modified)**: `app/main.py` lifespan wiring (build/start/stop the new scheduler); `app/core/config/settings.py` (refresh interval setting only); account-summary mappers in `app/modules/accounts/` and the dashboard mapper to join the cached snapshot onto `AccountSummary`. +- **Frontend (new/modified)**: `features/accounts/schemas.ts` + `features/dashboard/schemas.ts` (two new fields); `features/accounts/api.ts` (consume client); `features/accounts/sorting.ts` (new sort mode); `features/accounts/components/account-actions.tsx` (Reset button); `features/accounts/components/account-list-item.tsx` (count badge); `features/dashboard/components/account-list.tsx` + `account-card.tsx` (Reset button); `components/layout/app-header.tsx` (Accounts total badge); `utils/formatters.ts` (single-unit countdown + local expiry timestamp formatter); reuse of `components/confirm-dialog.tsx` + `hooks/use-dialog-state.ts` for the confirmation flow. +- **Upstream contract**: undocumented OpenAI endpoints under `https://chatgpt.com/backend-api/wham/rate-limit-reset-credits`; behavior is best-effort and may change upstream (see context doc). +- **In-memory only**: no DB schema migration; snapshots are lost on restart and repopulated on the next tick. +- **Tests**: pytest for client/store/scheduler/API/mapper; vitest (or equivalent) for formatter boundaries, badge cap, button visibility, confirm flow, and sort comparator. diff --git a/openspec/changes/add-rate-limit-reset-credits/specs/api-keys/spec.md b/openspec/changes/add-rate-limit-reset-credits/specs/api-keys/spec.md new file mode 100644 index 000000000..2d81d7a4c --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/specs/api-keys/spec.md @@ -0,0 +1,52 @@ +## ADDED Requirements + +### Requirement: API keys can inspect and redeem reset credits within their account pool + +The system SHALL expose `GET /v1/reset-credit` and `POST /v1/reset-credit` for API-key-authenticated self-service reset-credit access. Both routes MUST require a valid `Authorization: Bearer sk-clb-...` header even when `api_key_auth_enabled` is false globally. Validation failures MUST use the existing OpenAI error envelope used by `/v1/*` routes. + +The target account pool SHALL be derived from the authenticated API key. If `account_assignment_scope_enabled=true`, only `assigned_account_ids` SHALL be eligible. If account scope is not enabled, all selectable accounts SHALL be eligible. + +`GET /v1/reset-credit` SHALL return only credits for the authenticated key's eligible account pool. `POST /v1/reset-credit` SHALL reject requests whose `account_id` is outside that pool. + +On a successful `POST /v1/reset-credit` redemption, the system SHALL invalidate the redeemed account's cached reset-credit snapshot, force a usage refresh for that account, and invalidate account-selection cache state when that usage refresh writes updated usage. A failed or empty post-redeem usage refresh SHALL NOT roll back the successful credit redemption response. + +#### Scenario: Missing API key is rejected + +- **WHEN** a client calls `GET /v1/reset-credit` or `POST /v1/reset-credit` without a Bearer token +- **THEN** the system returns 401 in the OpenAI error format + +#### Scenario: Invalid API key is rejected + +- **WHEN** a client calls `GET /v1/reset-credit` or `POST /v1/reset-credit` with an unknown, expired, or inactive Bearer key +- **THEN** the system returns 401 in the OpenAI error format + +#### Scenario: Scoped API key sees only assigned accounts + +- **WHEN** an API key has account scope enabled with assigned accounts +- **AND** the client calls `GET /v1/reset-credit` +- **THEN** the response includes reset-credit entries only for those assigned accounts + +#### Scenario: Unscoped API key can read the full selectable pool + +- **WHEN** an API key has account scope disabled +- **AND** the client calls `GET /v1/reset-credit` +- **THEN** the response may include reset-credit entries for any selectable account that currently has an available cached credit + +#### Scenario: Out-of-pool account is rejected on redeem + +- **WHEN** a client calls `POST /v1/reset-credit` with an `account_id` outside the authenticated API key's eligible pool +- **THEN** the system returns 403 without redeeming any credit + +#### Scenario: Self-service reset-credit works while global proxy auth is disabled + +- **WHEN** `api_key_auth_enabled` is false and a client calls `GET /v1/reset-credit` or `POST /v1/reset-credit` with a valid Bearer key +- **THEN** the system still authenticates that key and applies the same account-pool rules + +#### Scenario: Successful self-service redemption refreshes usage for immediate follow-up traffic + +- **GIVEN** an eligible account has a redeemable reset credit and persisted usage/account state that still reflects a blocked window +- **WHEN** a client successfully calls `POST /v1/reset-credit` for that account +- **THEN** the redeemed account's cached reset-credit snapshot is invalidated +- **AND** codex-lb forces a usage refresh for that account before returning +- **AND** any account-selection cache entry derived from the stale usage state is invalidated when the refresh writes updated usage +- **AND** the response still returns the upstream `{code, windows_reset, redeemed_at}` success payload diff --git a/openspec/changes/add-rate-limit-reset-credits/specs/frontend-architecture/spec.md b/openspec/changes/add-rate-limit-reset-credits/specs/frontend-architecture/spec.md new file mode 100644 index 000000000..1ae91f6dc --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/specs/frontend-architecture/spec.md @@ -0,0 +1,109 @@ +## ADDED Requirements + +### Requirement: Accounts page exposes a reset-credits redeem action + +The Accounts page per-account action bar SHALL render a `Reset (N)` button next to the existing Export button with matching button styling whenever the account reports `available_reset_credits > 0`, where `N` is the available reset-credit count for that account. The button SHALL be hidden when `available_reset_credits` is `0`. Activating the button SHALL open a confirmation dialog that describes redeeming the soonest-expiring banked reset credit for that account and, when credit details are available, shows the soonest credit's expiry in local time using `YYYY-MM-DD HH:MM:SS`. Confirming SHALL submit a redeem request for that account and refresh account data on success. + +#### Scenario: Reset button mirrors Export styling and placement +- **WHEN** the Accounts page renders the per-account action bar for an account with `available_reset_credits > 0` +- **THEN** a `Reset (N)` button appears immediately next to the Export button +- **AND** the button uses the same size, variant, and class as the Export button + +#### Scenario: Reset button hidden when no credits available +- **WHEN** an account reports `available_reset_credits: 0` +- **THEN** the per-account action bar renders no "Reset" button + +#### Scenario: Confirmation required before redeem +- **WHEN** the operator clicks the "Reset" button +- **THEN** a confirmation dialog opens describing the soonest-expiring banked reset-credit redeem action +- **AND** no redeem request is sent until the operator confirms + +#### Scenario: Confirmation dialog shows local expiry timestamp +- **WHEN** the operator opens the reset-credit confirmation dialog and credit details include an expiry timestamp +- **THEN** the dialog renders the credit expiry in local time using `YYYY-MM-DD HH:MM:SS` + +### Requirement: AccountListItem displays a reset-credits count badge + +The Accounts page `AccountListItem` SHALL render a count badge pinned to the right-upper radius of the item whenever the account reports `available_reset_credits > 0`. The badge SHALL display the integer count, capped visually at `"99+"` when the count exceeds 99. The badge SHALL be absent when `available_reset_credits` is `0`. + +#### Scenario: Badge shows the available count +- **WHEN** an `AccountListItem` renders for an account with `available_reset_credits: 3` +- **THEN** a count badge pinned to the item's right-upper radius displays `3` + +#### Scenario: Badge caps at 99+ +- **WHEN** an `AccountListItem` renders for an account with `available_reset_credits: 120` +- **THEN** the count badge displays `99+` + +#### Scenario: Badge absent when zero +- **WHEN** an `AccountListItem` renders for an account with `available_reset_credits: 0` +- **THEN** no count badge is rendered + +### Requirement: Accounts page can sort by available reset credits + +The Accounts page sort selector SHALL offer a "Most reset credits" option and SHALL use it as the default Accounts page ordering. That ordering SHALL sort accounts by `available_reset_credits` descending. Ties SHALL be broken by `reset_credit_nearest_expires_at` ascending (soonest expiring first), and accounts with no expiry SHALL sort after accounts that have one. + +#### Scenario: More available credits sorts first +- **WHEN** the operator opens the Accounts page with the default sort mode +- **AND** account A has `available_reset_credits: 4` and account B has `available_reset_credits: 1` +- **THEN** account A appears before account B + +#### Scenario: Tie breaks by soonest expiry +- **WHEN** two accounts have equal `available_reset_credits` +- **AND** one account's soonest credit expires before the other's +- **THEN** the account with the earlier `reset_credit_nearest_expires_at` appears first + +### Requirement: Dashboard accounts section exposes a reset-credits redeem action + +The Dashboard Accounts section SHALL render a reset action next to the existing Details action in both the table and grid views for any account with `available_reset_credits > 0`. The grid view label SHALL read `Reset (N)`. The table view MAY remain icon-only, but its tooltip/title SHALL include the available reset-credit count. The action SHALL be absent when `available_reset_credits` is `0`. Activating the action SHALL open the same confirmation flow as the Accounts page reset action. + +#### Scenario: Table view shows reset next to details +- **WHEN** the Dashboard Accounts section renders in table view for an account with `available_reset_credits > 0` +- **THEN** a "Reset" action appears in the same action cell as the Details action + +#### Scenario: Grid view shows reset next to details +- **WHEN** the Dashboard Accounts section renders in grid view for an account with `available_reset_credits > 0` +- **THEN** a `Reset (N)` button appears next to the Details button on the account card + +#### Scenario: Reset action absent when no credits +- **WHEN** an account reports `available_reset_credits: 0` +- **THEN** the Dashboard Accounts section renders no "Reset" action for that account in either view + +### Requirement: Dashboard header shows the total available reset-credit count + +The dashboard top navigation SHALL render the total available reset-credit count on the Accounts tab, pinned to the tab's upper-right radius. The total SHALL equal the sum of `available_reset_credits` across the current account list data. The badge SHALL display `99+` when the total exceeds 99 and SHALL be hidden when the total is 0. + +#### Scenario: Accounts tab shows the summed total +- **WHEN** the current account list totals `available_reset_credits` to `14` +- **THEN** the Accounts nav tab displays a badge with `14` + +#### Scenario: Accounts tab caps large totals +- **WHEN** the current account list totals `available_reset_credits` to `120` +- **THEN** the Accounts nav tab displays a badge with `99+` + +#### Scenario: Accounts tab hides empty totals +- **WHEN** every account reports `available_reset_credits: 0` +- **THEN** the Accounts nav tab displays no reset-credit badge + +### Requirement: Reset actions display a single-unit expiry countdown + +Every "Reset" button SHALL display a small countdown label of the soonest-expiring credit's expiry, formatted as a single time unit: `"${d}d"` for any remaining duration of one day or more, `"${h}h"` for durations under one day but at least one hour, `"${m}m"` for durations under one hour but at least one minute, and `"now"` for durations under one minute. The label SHALL render in the destructive/red color when the remaining duration is strictly less than 7 days, and in the default muted color otherwise. + +#### Scenario: Days format for duration at or above one day +- **WHEN** a Reset button renders for a credit whose `expires_at` is 12 days away +- **THEN** the countdown label reads `12d` +- **AND** the label uses the default muted color + +#### Scenario: Red color under seven days +- **WHEN** a Reset button renders for a credit whose `expires_at` is 6 days away +- **THEN** the countdown label reads `6d` +- **AND** the label uses the destructive/red color + +#### Scenario: Hours and minutes use the smaller unit +- **WHEN** a Reset button renders for a credit whose `expires_at` is 13 hours away +- **THEN** the countdown label reads `13h` +- **AND** the label uses the destructive/red color + +#### Scenario: Sub-minute duration shows now +- **WHEN** a Reset button renders for a credit whose `expires_at` is 30 seconds away +- **THEN** the countdown label reads `now` +- **AND** the label uses the destructive/red color diff --git a/openspec/changes/add-rate-limit-reset-credits/specs/rate-limit-reset-credits/context.md b/openspec/changes/add-rate-limit-reset-credits/specs/rate-limit-reset-credits/context.md new file mode 100644 index 000000000..32f329c65 --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/specs/rate-limit-reset-credits/context.md @@ -0,0 +1,128 @@ +# Rate-Limit Reset Credits Context + +## Purpose + +codex-lb polls OpenAI's banked ("savable") rate-limit reset credits per account, caches them +in memory, and lets dashboard operators redeem the soonest-expiring credit for any account +without leaving the dashboard. The credit is a ChatGPT-subscription entitlement granted by +OpenAI; codex-lb is spending a credit OpenAI already gave the account — it does not bypass +any rate limit. + +## Upstream Source + +The credits endpoints live under `https://chatgpt.com/backend-api/wham`: + +| Endpoint | Method | Purpose | +|----------|--------|---------| +| `/wham/rate-limit-reset-credits` | GET | List banked credits + `available_count` | +| `/wham/rate-limit-reset-credits/consume` | POST | Redeem one credit (body: `credit_id`, `redeem_request_id`) | + +Both require `Authorization: Bearer ` and `chatgpt-account-id: ` +headers. The consume body returns `{code, credit: {id, status, redeemed_at, ...}, windows_reset}`. + +These endpoints are undocumented and were reverse-engineered from the official +`openai.chatgpt` VS Code extension's webview bundle. The canonical external reference is +[`aaamosh/codex-reset`](https://github.com/aaamosh/codex-reset) — a single-account CLI +implementation that codex-lb's multi-account, dashboard-driven, in-memory-cached variant +is based on. OpenAI may rename, gate, or remove these endpoints at any time; the codex-lb +client treats non-200, non-JSON, and schema-drifted 200 responses defensively. + +## Decisions + +- **In-memory only.** No DB column, no migration. Each replica refreshes its own process-local + snapshots, which repopulate within one tick of startup. Restart cost: up to 60s of + `available_reset_credits: 0` on that replica. +- **Server picks the credit, not the client.** `POST /consume` takes only the account id; + the server selects the soonest-expiring available credit from the freshest snapshot and + generates the `redeem_request_id`. Avoids stale-UI and clock-skew races. +- **Never mutates account status.** Account status is owned by usage refresh + (see `usage-refresh-policy`). Reset-credit polling failure logs and retains the prior + snapshot; it does not deactivate, rate-limit, or quota-block any account. +- **Dedicated scheduler, not folded into usage refresh.** Reuses the `UsageRefreshScheduler` + loop shape (`asyncio.Lock`-guarded, configurable cadence) but intentionally does not use + leader election because the cache is process-local. The scheduler always starts with the + app; only the interval is configurable. See `design.md` for the rationale. + +## Failure Modes + +- **Upstream returns 200 but the rate-limit window doesn't move.** Per upstream behavior + the credit is still consumed. The dashboard requires explicit confirmation before + redeeming; on success we invalidate the cache and let the next tick reconcile + `available_count`. +- **Snapshot is empty/stale.** UI hides all reset affordances for that account + (`available_reset_credits: 0`). Not an error — wait one tick. +- **Fresh consume preflight disproves a cached credit.** If the live pre-consume fetch says + `available_count: 0` or returns no available items, codex-lb overwrites the cached snapshot + with that fresh upstream state before returning `409`, so the dashboard does not keep + advertising a stale `Reset (N)` action until the next scheduler tick. +- **Account becomes ineligible after a successful snapshot.** Scheduler skips paused, + reauth-required, deactivated, or account-id-less accounts, so dashboard reads also check + current account eligibility before serving cached reset credits. If the account is + ineligible, the read returns no snapshot and invalidates the stale cache entry. +- **Upstream 401/403/auth-expired.** Logged; prior snapshot retained. Does NOT deactivate + the account. If the token is genuinely expired, usage refresh / OAuth refresh owns the + deactivation path. +- **Concurrent consume clicks.** Redemption is serialized per account so two overlapping + consume requests cannot forward the same cached `credit_id` upstream. After the first + request finishes, the second request re-reads the account snapshot and either sees a + refreshed state or fails with a dashboard conflict when no credit is still available. +- **Successful self-service redeem leaves stale usage state behind.** The `/v1/reset-credit` + success path force-refreshes usage for the redeemed account and invalidates the + load-balancer's account-selection cache when that refresh writes updated usage, so + `rate_limited` / `quota_exceeded` recovery can take effect for immediate follow-up + `/v1/*` traffic instead of waiting for the next periodic usage tick. +- **Upstream consume failures.** Client-facing upstream failures are preserved as dashboard + errors (`401`, `403`, `409`), while other consume failures surface as dashboard `503` + responses instead of falling into the generic internal-error handler. + +## Example: list response + +```json +{ + "credits": [ + { + "id": "RateLimitResetCredit_test", + "reset_type": "codex_rate_limits", + "status": "available", + "granted_at": "2026-06-12T01:29:41.346025Z", + "expires_at": "2026-07-12T01:29:41.346025Z", + "redeem_started_at": null, + "redeemed_at": null, + "profile_image_url": "https://openaiassets.blob.core.windows.net/$web/codex/codex-icon-200.png", + "profile_user_id": "Codex Team", + "title": "One free rate limit reset", + "description": "Thanks for using Codex! You've been granted one free rate limit reset." + } + ], + "available_count": 1 +} +``` + +## Example: consume response + +```json +{ + "code": "reset", + "credit": { + "id": "RateLimitResetCredit_...", + "reset_type": "codex_rate_limits", + "status": "redeemed", + "redeemed_at": "2026-06-13T13:12:31Z" + }, + "windows_reset": 1 +} +``` + +## Operational Notes + +- The 60s cadence matches usage refresh, but each replica polls because each replica serves + dashboard reads from its own process-local snapshot cache. +- A credit is consumed as soon as upstream returns 200 — treat the confirmation dialog as + the point of no return. + +## Related Work + +- Reference CLI: [`aaamosh/codex-reset`](https://github.com/aaamosh/codex-reset) +- Sibling capability: [`usage-refresh-policy`](../../specs/usage-refresh-policy/) — owns + account-status derivation and the `/wham/usage` 60s polling pattern this mirrors +- OpenAI announcement: [Flexible rate-limit resets for Codex](https://community.openai.com/t/flexible-rate-limit-resets-for-codex/1383470) diff --git a/openspec/changes/add-rate-limit-reset-credits/specs/rate-limit-reset-credits/spec.md b/openspec/changes/add-rate-limit-reset-credits/specs/rate-limit-reset-credits/spec.md new file mode 100644 index 000000000..205128c80 --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/specs/rate-limit-reset-credits/spec.md @@ -0,0 +1,125 @@ +## ADDED Requirements + +### Requirement: Reset credits are polled per account on a fixed cadence + +The system SHALL poll upstream `GET /wham/rate-limit-reset-credits` for each eligible account on a configurable cadence that defaults to 60 seconds, using that account's stored OAuth bearer token and `chatgpt-account-id`. The scheduler SHALL always start with the application lifespan. Because snapshots are kept in process-local memory, every running replica SHALL refresh its own snapshot cache instead of relying on leader election. The poll SHALL skip any account that is paused, requires reauthentication, deactivated, or lacks a usable `chatgpt-account-id`. + +#### Scenario: Default cadence polls every 60 seconds +- **WHEN** the application starts with default settings +- **THEN** each eligible account's credits are fetched from upstream at most once per 60 seconds + +#### Scenario: Every replica refreshes its local cache +- **WHEN** the application is deployed with multiple running replicas +- **THEN** each replica refreshes its own in-memory reset-credit snapshots on the configured cadence +- **AND** dashboard reads served by any replica can observe populated reset-credit data after that replica's refresh tick + +#### Scenario: Ineligible accounts are skipped +- **WHEN** an account is persisted as `paused`, `reauth_required`, or `deactivated` +- **THEN** the scheduler performs no upstream reset-credits fetch for that account +- **AND** the cached snapshot for that account (if any) is left untouched by the skip + +### Requirement: Reset credit snapshots are cached in memory keyed by account + +The system SHALL store the most recent successful reset-credits response per account in an in-memory store keyed by account id. The store SHALL be concurrency-safe and SHALL provide an `invalidate(account_id)` operation. Account-summary mappers SHALL join the cached snapshot onto each account summary, exposing `available_reset_credits` (integer) and `reset_credit_nearest_expires_at` (ISO timestamp or null). Accounts with no cached snapshot SHALL expose `available_reset_credits: 0` and `reset_credit_nearest_expires_at: null`. + +#### Scenario: Account summary reflects cached credits +- **GIVEN** an account has a cached reset-credits snapshot with `available_count: 2` and a soonest expiry of `2026-07-10T00:00:00Z` +- **WHEN** the account-summary mapper builds the summary for that account +- **THEN** the summary exposes `available_reset_credits: 2` and `reset_credit_nearest_expires_at: "2026-07-10T00:00:00Z"` + +#### Scenario: Missing cache presents as zero credits +- **GIVEN** an account has no cached reset-credits snapshot (e.g. immediately after restart) +- **WHEN** the account-summary mapper builds the summary for that account +- **THEN** the summary exposes `available_reset_credits: 0` and `reset_credit_nearest_expires_at: null` + +#### Scenario: Invalidate forces re-fetch on next tick +- **WHEN** a caller invokes `invalidate(account_id)` for an account +- **THEN** subsequent reads for that account return no cached snapshot +- **AND** the next scheduler tick fetches a fresh snapshot from upstream + +#### Scenario: In-flight refresh cannot restore an invalidated snapshot +- **GIVEN** a scheduler refresh starts fetching reset credits for an account +- **AND** another caller invokes `invalidate(account_id)` before that refresh stores its fetched response +- **WHEN** the refresh completes +- **THEN** the stale fetched response MUST NOT be written back into the cache + +#### Scenario: Dashboard read invalidates stale snapshots for ineligible accounts +- **GIVEN** an account has a cached reset-credits snapshot +- **AND** the account is now persisted as `paused`, `reauth_required`, `deactivated`, or no longer has a usable `chatgpt-account-id` +- **WHEN** the dashboard invokes `GET /api/accounts/{id}/rate-limit-reset-credits` +- **THEN** the endpoint returns `null` without calling upstream +- **AND** the cached snapshot for that account is invalidated + +### Requirement: Operators can redeem the soonest-expiring available credit + +The system SHALL expose a dashboard endpoint `POST /api/accounts/{account_id}/rate-limit-reset-credits/consume` that redeems exactly one credit for the named account. The endpoint SHALL select, from the freshest cached snapshot, the credit whose `status` is `available` with the smallest `expires_at`, generate a `redeem_request_id` (UUID v4), and forward `{credit_id, redeem_request_id}` to upstream `POST /wham/rate-limit-reset-credits/consume` using the account's bearer token and `chatgpt-account-id`. A cached snapshot with `available_count <= 0` MUST be treated as having no redeemable credits, even if the cached `credits` list contains an item marked `available`. When the fresh pre-consume fetch reports `available_count <= 0` or no available credit items, the endpoint SHALL replace any prior cached snapshot for that account with the fresh upstream snapshot before returning a conflict. On a 200 response the endpoint SHALL invalidate the cached snapshot for that account and return `{code, windows_reset, redeemed_at}`. The endpoint SHALL require dashboard write access; read-only guests MUST be refused. + +#### Scenario: Consume selects the soonest-expiring credit +- **GIVEN** an account has cached credits with expiries `2026-07-10Z` and `2026-06-20Z`, both `status: available` +- **WHEN** the operator invokes `POST /api/accounts/{id}/rate-limit-reset-credits/consume` +- **THEN** the request forwarded to upstream carries the `credit_id` whose `expires_at` is `2026-06-20Z` + +#### Scenario: Successful consume invalidates the cache +- **GIVEN** the operator invokes consume for an account with at least one available credit +- **WHEN** upstream returns `200` with `{code: "reset", windows_reset: 1, credit: {...}}` +- **THEN** the cached snapshot for that account is invalidated +- **AND** the response returned to the dashboard is `{code, windows_reset, redeemed_at}` derived from the upstream response + +#### Scenario: Concurrent consume requests for one account are serialized +- **GIVEN** two operators invoke `POST /api/accounts/{id}/rate-limit-reset-credits/consume` at nearly the same time for the same account +- **WHEN** the first request is still redeeming a credit +- **THEN** the second request MUST wait for the first request to finish before re-reading that account's cached snapshot +- **AND** the same cached `credit_id` MUST NOT be sent to upstream twice by those concurrent requests + +#### Scenario: Upstream consume failures surface as dashboard errors +- **GIVEN** an operator invokes `POST /api/accounts/{id}/rate-limit-reset-credits/consume` +- **WHEN** upstream returns `401`, `403`, or `409` +- **THEN** the dashboard endpoint returns the same client-facing status class instead of a generic `500` +- **AND** other upstream consume failures return a dashboard `503` + +#### Scenario: Read-only guests cannot redeem +- **GIVEN** a dashboard session authenticated as a read-only guest +- **WHEN** the guest invokes `POST /api/accounts/{id}/rate-limit-reset-credits/consume` +- **THEN** the request is refused before any upstream call is made + +#### Scenario: Consume with no available credit returns a client error +- **GIVEN** an account whose cached snapshot reports `available_count: 0` (or has no snapshot) +- **WHEN** the operator invokes `POST /api/accounts/{id}/rate-limit-reset-credits/consume` +- **THEN** the endpoint returns a `409` (or equivalent client-error) without calling upstream + +#### Scenario: Fresh empty consume fetch replaces a stale cached snapshot +- **GIVEN** an account has a cached reset-credits snapshot showing at least one available credit +- **AND** the fresh pre-consume upstream fetch returns `available_count: 0` or no `status: available` items +- **WHEN** the operator invokes `POST /api/accounts/{id}/rate-limit-reset-credits/consume` +- **THEN** the endpoint returns a `409` (or equivalent client-error) +- **AND** the cached snapshot for that account is replaced with the fresh upstream snapshot before the response is returned + +### Requirement: Reset credit polling failure does not mutate account status + +The reset-credits refresh scheduler SHALL NOT transition any account's persisted status (`active`, `rate_limited`, `quota_exceeded`, `paused`, `deactivated`) in response to upstream reset-credits responses. On upstream error (non-200, non-JSON, malformed 200 payload, network, or auth-like failure) the scheduler SHALL log the failure and either keep the prior cached snapshot or leave the cache unset; it SHALL NOT propagate the failure to account-status derivation. + +#### Scenario: Upstream 401 on reset-credits does not deactivate the account +- **WHEN** the scheduler receives an HTTP `401` from `GET /wham/rate-limit-reset-credits` for an account +- **THEN** the account's persisted status is unchanged +- **AND** any prior cached snapshot for that account is retained + +#### Scenario: Upstream 5xx retains the prior snapshot +- **GIVEN** an account has a cached snapshot from a prior successful tick +- **WHEN** the scheduler receives an HTTP `503` on the next reset-credits tick +- **THEN** the cached snapshot is retained +- **AND** the failure is logged + +#### Scenario: Malformed 200 response is not cached as success +- **GIVEN** an account has a cached snapshot from a prior successful tick +- **WHEN** upstream returns HTTP `200` with a non-object body or a body missing required reset-credit fields +- **THEN** the response is treated as an upstream failure +- **AND** the cached snapshot is retained + +### Requirement: Reset credit polling interval is configurable + +The system SHALL expose setting `rate_limit_reset_credits_refresh_interval_seconds` (default `60`) to control the polling cadence. The system SHALL NOT expose a separate enable/disable toggle for reset-credit polling. + +#### Scenario: Operator tunes the polling interval +- **GIVEN** `rate_limit_reset_credits_refresh_interval_seconds` is set to `120` +- **WHEN** the application starts and runs +- **THEN** each eligible account's credits are fetched from upstream at most once per 120 seconds diff --git a/openspec/changes/add-rate-limit-reset-credits/tasks.md b/openspec/changes/add-rate-limit-reset-credits/tasks.md new file mode 100644 index 000000000..2fd3b6cf8 --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/tasks.md @@ -0,0 +1,52 @@ +## 1. Backend foundation (settings, upstream client, in-memory store) + +- [x] 1.1 Add setting `rate_limit_reset_credits_refresh_interval_seconds` (default `60`) to `app/core/config/settings.py`; reset-credit polling itself is always on +- [x] 1.2 Create `app/core/clients/rate_limit_reset_credits.py` mirroring `app/core/clients/usage.py`: `fetch_reset_credits(access_token, account_id, *, base_url, timeout)` → GET `/wham/rate-limit-reset-credits`, and `consume_reset_credit(access_token, account_id, credit_id, *, base_url, timeout)` → POST `/wham/rate-limit-reset-credits/consume` with body `{credit_id, redeem_request_id: uuid4()}`. Reuse the same header-construction rules (skip `chatgpt-account-id` for `email_`/`local_` prefixes) and base-url normalization +- [x] 1.3 Define pydantic models for the upstream payloads: `ResetCreditItem` (id, reset_type, status, granted_at, expires_at, title, description, redeem_started_at, redeemed_at), `ResetCreditsResponse` (credits: list, available_count: int), `ConsumeResetCreditResponse` (code, credit, windows_reset) +- [x] 1.4 Create `app/modules/rate_limit_reset_credits/store.py` mirroring `app/modules/proxy/rate_limit_cache.py`: `RateLimitResetCreditsStore` with `anyio.Lock`-guarded `set(account_id, snapshot)`, `get(account_id) -> Snapshot | None`, `invalidate(account_id=None)`. Snapshot exposes `available_count`, `nearest_expires_at`, and the items list. Expose a module-level singleton + `get_rate_limit_reset_credits_store()` accessor + +## 2. Backend scheduler, API, mapper, lifespan wiring + +- [x] 2.1 Create `app/core/usage/reset_credits_refresh_scheduler.py` mirroring `app/core/usage/refresh_scheduler.py`: `RateLimitResetCreditsRefreshScheduler` dataclass with `asyncio.Lock`-guarded `_refresh_once` that runs in every replica, lists accounts, skips paused/deactivated/missing-`chatgpt-account-id`, decrypts `access_token_encrypted`, calls `fetch_reset_credits`, and stores the snapshot. On upstream error: log + retain prior snapshot; do NOT mutate account status. Add `build_rate_limit_reset_credits_scheduler()` factory +- [x] 2.2 Wire the new scheduler into `app/main.py` lifespan alongside `usage_scheduler`: build (~line 148), start (~154), stop (~314) +- [x] 2.3 Create `app/modules/rate_limit_reset_credits/api.py` with `GET /api/accounts/{account_id}/rate-limit-reset-credits` (returns cached snapshot or `null`) and `POST /api/accounts/{account_id}/rate-limit-reset-credits/consume` (selects soonest-`expires_at` available credit from the freshest snapshot, generates `redeem_request_id`, calls upstream, invalidates the cached snapshot, returns `{code, windows_reset, redeemed_at}`). Use `validate_dashboard_session` for GET and `require_dashboard_write_access` for POST. Return `409` when no credit is available. Register the router in `app/main.py` +- [x] 2.4 Extend the AccountSummary mapper(s) in `app/modules/accounts/` and the dashboard mapper to join the cached snapshot onto each returned account: add `available_reset_credits: int` (0 when no snapshot) and `reset_credit_nearest_expires_at: datetime | None` (null when no snapshot) +- [x] 2.5 Update the backend pydantic response schemas (`AccountSummary` / equivalent) to declare the two new fields + +## 3. Frontend schemas, API client, formatter + +- [x] 3.1 Add `availableResetCredits: number` and `resetCreditNearestExpiresAt: string | null` to `AccountSummary` in both `frontend/src/features/accounts/schemas.ts` and `frontend/src/features/dashboard/schemas.ts` +- [x] 3.2 Add `consumeRateLimitResetCredit(accountId): Promise<{ code: string; windowsReset: number; redeemedAt: string }>` to `frontend/src/features/accounts/api.ts` posting to `/api/accounts/{id}/rate-limit-reset-credits/consume`. On success, invalidate the `['accounts']` and `['dashboard']` TanStack Query keys +- [x] 3.3 Add `formatSingleUnitRemaining(expiresAtIso: string): { label: string; expiringSoon: boolean }` to `frontend/src/utils/formatters.ts`: `"${d}d"` for ≥1 day, `"${h}h"` for ≥1 hour, `"${m}m"` for ≥1 minute, `"now"` otherwise; `expiringSoon = ms < 7 * 86_400_000`. Sit it next to the existing `formatResetRelative`/`formatQuotaResetLabel` helpers + +## 4. Frontend UI components + +- [x] 4.1 Add the count badge to `frontend/src/features/accounts/components/account-list-item.tsx`: an absolutely-positioned circle on the right-upper radius showing the integer count or `"99+"` when `> 99`. Render only when `availableResetCredits > 0` +- [x] 4.2 Add the `Reset (N)` button to `frontend/src/features/accounts/components/account-actions.tsx` immediately after the Export button, matching its `size="sm" variant="outline" className="h-8 gap-1.5 text-xs"` style, with a `RotateCcw` icon, a single-unit countdown label (using 3.3) placed at the button's right-upper radius, and destructive/red label color when `expiringSoon`. Render only when `availableResetCredits > 0`. Wire `onClick` to open the confirmation dialog +- [x] 4.3 Add a reset action to `frontend/src/features/dashboard/components/account-list.tsx` (table view) inside the existing Details action cell, matching the `h-7 w-7` icon-button style with the countdown and count exposed in the `title` tooltip. Render only when `availableResetCredits > 0` +- [x] 4.4 Add a `Reset (N)` button to `frontend/src/features/dashboard/components/account-card.tsx` (grid view) next to the Details button, matching the `h-7 gap-1.5` text style with the single-unit countdown label. Render only when `availableResetCredits > 0` +- [x] 4.5 Implement the confirmation dialog (reuse `frontend/src/components/confirm-dialog.tsx` + `frontend/src/hooks/use-dialog-state.ts`, same shape as the delete-account dialog): body describes the soonest-expiring banked reset-credit redeem action and shows `expires_at` formatted as local `YYYY-MM-DD HH:MM:SS` when credit details are available. On confirm → call `consumeRateLimitResetCredit(accountId)` → success/failure toast → query invalidation +- [x] 4.6 Add a "Most reset credits" option to the Accounts page sort selector in `frontend/src/features/accounts/sorting.ts` and make it the default Accounts page sort mode: comparator orders by `availableResetCredits` desc, tiebreak by `resetCreditNearestExpiresAt` asc (soonest first), accounts with null expiry last. Add the localized dropdown label +- [x] 4.7 Add a summed reset-credit badge to `frontend/src/components/layout/app-header.tsx` for the Accounts nav tab, capped at `99+` + +## 5. Tests + +- [x] 5.1 Backend — `app/core/clients/rate_limit_reset_credits.py`: header construction (account-id skip rule), base-url normalization, consume body shape, JSON parse on 200, error handling on non-200/non-JSON +- [x] 5.2 Backend — `app/modules/rate_limit_reset_credits/store.py`: `set`/`get`/`invalidate` (single + all), concurrency under `anyio.Lock`, missing-account returns `None` +- [x] 5.3 Backend — `reset_credits_refresh_scheduler.py`: every replica refreshes its local cache, paused/deactivated account skip, one-account failure doesn't break the loop, upstream error retains prior snapshot, account status is never mutated +- [x] 5.4 Backend — `rate_limit_reset_credits/api.py`: GET returns cached snapshot / `null` on miss; POST selects soonest expiry, calls upstream with fresh `redeem_request_id`, invalidates cache, returns `{code, windows_reset, redeemed_at}`; write-access gating refuses guests; `409` when no available credit +- [x] 5.5 Backend — AccountSummary mapper: exposes the two new fields from a cached snapshot, returns `0`/`null` when no snapshot, does not crash when store is empty +- [x] 5.6 Frontend — `formatSingleUnitRemaining`: boundaries at 7d (color flip), 1d, 1h, 1m, and `now`; sub-minute and past timestamps both yield `"now"` +- [x] 5.7 Frontend — `AccountListItem` badge: renders count, `"99+"` at 100+, absent at 0 +- [x] 5.8 Frontend — Reset button visibility: rendered when `availableResetCredits > 0`, absent at 0, in all three surfaces (account-actions, dashboard table, dashboard grid) +- [x] 5.9 Frontend — confirm dialog → consume: confirmation calls `consumeRateLimitResetCredit`, shows the expiry in local `YYYY-MM-DD HH:MM:SS`, success path invalidates queries, failure path surfaces a toast and does not invalidate +- [x] 5.10 Frontend — "Most reset credits" sort: comparator orders by count desc with soonest-expiry tiebreak, null-expiry accounts last +- [x] 5.11 Frontend — Accounts nav badge: shows the summed total, caps at `99+`, and hides at zero + +## 6. Validation and OpenSpec hygiene + +- [x] 6.1 Run `openspec validate add-rate-limit-reset-credits --strict` and resolve any findings +- [x] 6.2 Run `openspec validate --specs --strict` to confirm no main-spec drift +- [ ] 6.3 Run backend checks: `uv run ruff check && uv run ruff format --check && uv run pytest` (or the repo's documented equivalent) +- [x] 6.4 Run frontend checks: `pnpm -C frontend lint && pnpm -C frontend typecheck && pnpm -C frontend test` (or the repo's documented equivalent) +- [ ] 6.5 Manually verify the three Reset button placements, the per-button count labels, the Accounts-nav total badge cap behavior, the countdown color flip at 7d, the local expiry timestamp, the confirm flow, and the new sort option against the spec scenarios diff --git a/openspec/changes/add-rate-limit-reset-credits/verify-report.md b/openspec/changes/add-rate-limit-reset-credits/verify-report.md new file mode 100644 index 000000000..b3be3cc76 --- /dev/null +++ b/openspec/changes/add-rate-limit-reset-credits/verify-report.md @@ -0,0 +1,78 @@ +## Verification Report + +### `openspec validate add-rate-limit-reset-credits --strict` + +- Result: passed +- Output: `Change 'add-rate-limit-reset-credits' is valid` + +### `uv run pytest tests/integration/test_v1_reset_credit.py tests/unit/test_rate_limit_reset_credits_api.py -v` + +- Result: passed +- Output: `28 passed in 1.52s` + +### `uv run ruff check && uv run ruff format --check && uv run pytest` + +- First attempt: tool timeout at 120000 ms while `pytest` was still running +- Rerun with a longer timeout: passed +- `ruff check`: `All checks passed!` +- `ruff format --check`: `648 files already formatted` +- `pytest`: `3703 passed, 45 skipped, 3 warnings in 223.70s (0:03:43)` + +### Frontend Reset-Credit Contract Follow-Up + +- Focused red/green after tightening the frontend consume contract: + - `bun run test src/features/accounts/schemas.test.ts src/features/accounts/components/reset-credit-confirm-dialog.test.tsx` + - Red step failed as expected before the code change on schema strictness and top-level query invalidation expectations + - Green step passed with `12` tests passing +- Full frontend quality gate rerun after the fix: + - `bun run lint` passed + - `bun run typecheck` passed + - `bun run test` passed with `104` test files and `652` tests passing in `89.29s` +- Frontend test stderr still includes existing React `act(...)`, Recharts zero-size container, and jsdom `HTMLCanvasElement.getContext()` warnings during an otherwise passing run + +### Final Fresh Verification Snapshot + +- `openspec validate add-rate-limit-reset-credits --strict` passed +- `openspec instructions apply --change "add-rate-limit-reset-credits" --json` now reports `40/41` tasks complete with only `6.5` remaining +- `bun run lint && bun run typecheck && bun run test` passed again after the frontend contract fix + - `104` test files passed + - `652` tests passed + - Existing stderr warnings remained non-fatal + +### `openspec validate --specs --strict` + +- Result: passed +- Output: `Totals: 30 passed, 0 failed (30 items)` + +### Frontend verification + +- Repo package manager declaration: `frontend/package.json` declares `"packageManager": "bun@1.3.14"` +- Practical command form used in this worktree: `bun run lint`, `bun run typecheck`, and `bun run test` with `workdir=frontend` +- Note: `bun -C frontend ...` is not supported by the installed Bun CLI in this environment and fails with `error: Invalid Argument '-C'` + +#### `bun run lint` + +- Result: passed +- Output: `$ eslint .` + +#### `bun run typecheck` + +- Result: passed +- Output: `$ tsc -b` + +#### `bun run test` + +- Result: passed +- Output: `Test Files 104 passed (104)` +- Output: `Tests 652 passed (652)` +- Output: `Duration 91.73s` +- Notes: stderr included existing React `act(...)` warnings, Recharts zero-size container warnings, and jsdom `HTMLCanvasElement.getContext()` not-implemented warnings during the passing run + +### Manual Verification + +- Not performed in this implementation pass +- `openspec/changes/add-rate-limit-reset-credits/tasks.md` item `6.5` remains unchecked + +### Remaining OpenSpec Gaps + +- `6.5` remains unchecked because the requested manual UI verification was not performed in this implementation pass. diff --git a/openspec/changes/fix-v1-reset-credit-token-refresh/proposal.md b/openspec/changes/fix-v1-reset-credit-token-refresh/proposal.md new file mode 100644 index 000000000..a34e5ea44 --- /dev/null +++ b/openspec/changes/fix-v1-reset-credit-token-refresh/proposal.md @@ -0,0 +1,13 @@ +## Why + +`POST /v1/reset-credit` currently decrypts and forwards the persisted access token for the selected account without first refreshing it. When that stored token is expired or past the normal refresh threshold, self-service reset-credit redemption fails upstream with 401 until some unrelated traffic refreshes the account. + +## What Changes + +- Refresh the target account via `AuthManager.ensure_fresh` before decrypting the bearer token used by `POST /v1/reset-credit`. +- Keep the self-service redemption path aligned with the dashboard reset-credit consume flow so both surfaces redeem against fresh credentials. + +## Impact + +- API-key-authenticated self-service reset-credit redemption succeeds for accounts whose stored token is stale but refreshable. +- No request or response schema changes. diff --git a/openspec/changes/fix-v1-reset-credit-token-refresh/specs/api-keys/spec.md b/openspec/changes/fix-v1-reset-credit-token-refresh/specs/api-keys/spec.md new file mode 100644 index 000000000..d8538a439 --- /dev/null +++ b/openspec/changes/fix-v1-reset-credit-token-refresh/specs/api-keys/spec.md @@ -0,0 +1,31 @@ +## MODIFIED Requirements + +### Requirement: API keys can inspect and redeem reset credits within their account pool + +The system SHALL expose `GET /v1/reset-credit` and `POST /v1/reset-credit` for API-key-authenticated self-service reset-credit access. Both routes MUST require a valid `Authorization: Bearer sk-clb-...` header even when `api_key_auth_enabled` is false globally. Validation failures MUST use the existing OpenAI error envelope used by `/v1/*` routes. + +The target account pool SHALL be derived from the authenticated API key. If `account_assignment_scope_enabled=true`, only `assigned_account_ids` SHALL be eligible. If account scope is not enabled, all selectable accounts SHALL be eligible. + +`GET /v1/reset-credit` SHALL return only credits for the authenticated key's eligible account pool. `POST /v1/reset-credit` SHALL reject requests whose `account_id` is outside that pool. + +Before `POST /v1/reset-credit` decrypts and forwards the bearer token for the upstream consume call, the system SHALL refresh the target account with the normal account-token freshness rules and use the refreshed account credentials for the consume request. + +If that self-service credential refresh fails, `POST /v1/reset-credit` SHALL stop before the upstream consume call, return a client-actionable conflict response, and keep using the existing `/v1/*` OpenAI error envelope. + +On a successful `POST /v1/reset-credit` redemption, the system SHALL invalidate the redeemed account's cached reset-credit snapshot, force a usage refresh for that account, and invalidate account-selection cache state when that usage refresh writes updated usage. A failed or empty post-redeem usage refresh SHALL NOT roll back the successful credit redemption response. + +#### Scenario: Self-service redemption refreshes stale account credentials before consume + +- **GIVEN** an eligible account has a redeemable reset credit +- **AND** the persisted access token for that account is stale but refreshable +- **WHEN** a client successfully calls `POST /v1/reset-credit` for that account +- **THEN** codex-lb refreshes the account before decrypting the consume bearer token +- **AND** the upstream reset-credit consume call uses the refreshed account credentials + +#### Scenario: Self-service redemption surfaces refresh failures as conflicts + +- **GIVEN** an eligible account has a redeemable reset credit +- **AND** that account's credential refresh fails before the upstream consume call +- **WHEN** a client calls `POST /v1/reset-credit` for that account +- **THEN** codex-lb returns a conflict response in the standard `/v1/*` OpenAI error envelope +- **AND** codex-lb does not call upstream reset-credit consume for that request diff --git a/openspec/changes/fix-v1-reset-credit-token-refresh/tasks.md b/openspec/changes/fix-v1-reset-credit-token-refresh/tasks.md new file mode 100644 index 000000000..a08e0cee3 --- /dev/null +++ b/openspec/changes/fix-v1-reset-credit-token-refresh/tasks.md @@ -0,0 +1,4 @@ +- [x] Refresh the account in `POST /v1/reset-credit` before decrypting and forwarding the bearer token to upstream reset-credit consume. +- [x] Add regression coverage proving the self-service reset-credit path consumes with the refreshed token material. +- [x] Translate self-service reset-credit refresh failures into a client-actionable conflict response before upstream consume. +- [x] Validate OpenSpec after updating the delta spec. diff --git a/tests/conftest.py b/tests/conftest.py index bb7720f8d..8b9e57eb4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -28,6 +28,14 @@ from app.main import create_app # noqa: E402 +class _NoopScheduler: + async def start(self) -> None: + return None + + async def stop(self) -> None: + return None + + def _drop_test_migration_tables(sync_conn) -> None: sync_conn.execute(text("DROP TABLE IF EXISTS alembic_version")) sync_conn.execute(text("DROP TABLE IF EXISTS schema_migrations")) @@ -62,10 +70,18 @@ async def _noop_init_db() -> None: return None monkeypatch.setattr(main_module, "init_db", _noop_init_db) + monkeypatch.setattr(main_module, "build_rate_limit_reset_credits_scheduler", lambda: _NoopScheduler()) app = create_app() return app +@pytest.fixture(autouse=True) +def _disable_rate_limit_reset_credits_scheduler_startup(monkeypatch): + import app.main as main_module + + monkeypatch.setattr(main_module, "build_rate_limit_reset_credits_scheduler", lambda: _NoopScheduler()) + + @pytest_asyncio.fixture(scope="session", autouse=True) async def dispose_engine(): yield diff --git a/tests/integration/test_rate_limit_reset_credits_api.py b/tests/integration/test_rate_limit_reset_credits_api.py new file mode 100644 index 000000000..0b7d54fea --- /dev/null +++ b/tests/integration/test_rate_limit_reset_credits_api.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import base64 +import json +from datetime import datetime +from typing import Any + +import pytest + +from app.core.auth import generate_unique_account_id +from app.core.clients.rate_limit_reset_credits import ( + ConsumeResetCreditResponse, + ResetCreditItem, + ResetCreditsResponse, +) +from app.db.session import SessionLocal +from app.modules.rate_limit_reset_credits import api as reset_credits_api + +pytestmark = pytest.mark.integration + + +def _encode_jwt(payload: dict) -> str: + raw = json.dumps(payload, separators=(",", ":")).encode("utf-8") + body = base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + return f"header.{body}.sig" + + +async def _import_test_account(async_client, *, email: str, account_id: str) -> str: + payload = { + "email": email, + "chatgpt_account_id": account_id, + "https://api.openai.com/auth": {"chatgpt_plan_type": "plus"}, + } + auth_json = { + "tokens": { + "idToken": _encode_jwt(payload), + "accessToken": "access-token-not-a-real-secret", + "refreshToken": "refresh", + "accountId": account_id, + }, + } + files = {"auth_json": ("auth.json", json.dumps(auth_json), "application/json")} + response = await async_client.post("/api/accounts/import", files=files) + assert response.status_code == 200, response.text + return generate_unique_account_id(account_id, email) + + +def _credit(credit_id: str, *, expires_at: str = "2026-07-12T00:00:00Z") -> ResetCreditItem: + return ResetCreditItem.model_validate({"id": credit_id, "status": "available", "expires_at": expires_at}) + + +def _upstream_response(credits: list[ResetCreditItem], available_count: int | None = None) -> ResetCreditsResponse: + count = available_count if available_count is not None else len(credits) + return ResetCreditsResponse(credits=credits, available_count=count) + + +@pytest.mark.asyncio +async def test_consume_paused_account_returns_409(async_client, monkeypatch) -> None: + async def _should_not_fetch(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + raise AssertionError("paused account should not invoke upstream fetch") + + monkeypatch.setattr(reset_credits_api, "fetch_reset_credits", _should_not_fetch) + + account_id = await _import_test_account( + async_client, + email="reset-paused@example.com", + account_id="acc_reset_paused", + ) + pause_resp = await async_client.post(f"/api/accounts/{account_id}/pause") + assert pause_resp.status_code == 200 + + response = await async_client.post(f"/api/accounts/{account_id}/rate-limit-reset-credits/consume") + assert response.status_code == 409 + body = response.json() + assert body["error"]["code"] == "account_not_reset_credit_applicable" + + +@pytest.mark.asyncio +async def test_consume_active_account_returns_success_with_mocked_upstream(async_client, monkeypatch) -> None: + captured: dict[str, Any] = {} + + async def _fake_fetch(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + captured["fetch_account_id"] = account_id + captured["fetch_had_token"] = bool(access_token) + return _upstream_response([_credit("credit-1")]) + + async def _fake_consume( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + captured.update( + { + "consume_account_id": account_id, + "consume_credit_id": credit_id, + "consume_had_token": bool(access_token), + } + ) + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": { + "id": credit_id, + "status": "redeemed", + "redeemed_at": "2026-06-13T13:12:31Z", + }, + "windows_reset": 2, + } + ) + + async def _noop_refresh(account) -> None: # noqa: ANN001 + return None + + monkeypatch.setattr(reset_credits_api, "fetch_reset_credits", _fake_fetch) + monkeypatch.setattr(reset_credits_api, "consume_reset_credit", _fake_consume) + monkeypatch.setattr(reset_credits_api, "_build_refresh_usage_callback", lambda _context: _noop_refresh) + + account_id = await _import_test_account( + async_client, + email="reset-active@example.com", + account_id="acc_reset_active", + ) + + warm_response = await async_client.get(f"/api/accounts/{account_id}/rate-limit-reset-credits") + assert warm_response.status_code == 200, warm_response.text + + response = await async_client.post(f"/api/accounts/{account_id}/rate-limit-reset-credits/consume") + assert response.status_code == 200, response.text + body = response.json() + assert body["code"] == "reset" + assert body["windowsReset"] == 2 + assert body["redeemedAt"] is not None + assert datetime.fromisoformat(body["redeemedAt"].replace("Z", "+00:00")).year == 2026 + + assert captured["fetch_account_id"] == "acc_reset_active" + assert captured["fetch_had_token"] is True + assert captured["consume_account_id"] == "acc_reset_active" + assert captured["consume_credit_id"] == "credit-1" + assert captured["consume_had_token"] is True + + +@pytest.mark.asyncio +async def test_consume_without_cached_snapshot_returns_409_without_fetch(async_client, monkeypatch) -> None: + async def _should_not_fetch(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + raise AssertionError("uncached consume should not invoke upstream fetch") + + monkeypatch.setattr(reset_credits_api, "fetch_reset_credits", _should_not_fetch) + + account_id = await _import_test_account( + async_client, + email="reset-no-cache@example.com", + account_id="acc_reset_no_cache", + ) + + response = await async_client.post(f"/api/accounts/{account_id}/rate-limit-reset-credits/consume") + assert response.status_code == 409 + assert response.json()["error"]["code"] == "no_available_reset_credit" + + +@pytest.mark.asyncio +async def test_consume_reauth_required_account_returns_409(async_client, monkeypatch) -> None: + async def _should_not_fetch(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + raise AssertionError("reauth account should not invoke upstream fetch") + + monkeypatch.setattr(reset_credits_api, "fetch_reset_credits", _should_not_fetch) + + account_id = await _import_test_account( + async_client, + email="reset-reauth@example.com", + account_id="acc_reset_reauth", + ) + + async with SessionLocal() as session: + from sqlalchemy import update + + from app.db.models import Account, AccountStatus + + await session.execute( + update(Account).where(Account.id == account_id).values(status=AccountStatus.REAUTH_REQUIRED) + ) + await session.commit() + + response = await async_client.post(f"/api/accounts/{account_id}/rate-limit-reset-credits/consume") + assert response.status_code == 409 + assert response.json()["error"]["code"] == "account_not_reset_credit_applicable" + + +@pytest.mark.asyncio +async def test_get_populates_reset_credits_on_cache_miss(async_client, monkeypatch) -> None: + account_id = await _import_test_account( + async_client, + email="reset-get@example.com", + account_id="acc_reset_get", + ) + + async def _fake_fetch(access_token: str, chatgpt_account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + return _upstream_response( + [ + ResetCreditItem.model_validate( + { + "id": "credit-get", + "status": "available", + "expires_at": "2026-08-01T00:00:00Z", + "title": "Thanks for using Codex!", + } + ) + ] + ) + + monkeypatch.setattr(reset_credits_api, "fetch_reset_credits", _fake_fetch) + + response = await async_client.get(f"/api/accounts/{account_id}/rate-limit-reset-credits") + assert response.status_code == 200, response.text + body = response.json() + assert body["availableCount"] == 1 + assert body["credits"][0]["id"] == "credit-get" diff --git a/tests/integration/test_v1_reset_credit.py b/tests/integration/test_v1_reset_credit.py new file mode 100644 index 000000000..e1743b5b9 --- /dev/null +++ b/tests/integration/test_v1_reset_credit.py @@ -0,0 +1,1041 @@ +from __future__ import annotations + +import base64 +import json +from contextlib import asynccontextmanager +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from sqlalchemy import update + +from app.core.auth import generate_unique_account_id +from app.core.auth.refresh import RefreshError +from app.core.clients.rate_limit_reset_credits import ( + ConsumeResetCreditError, + ConsumeResetCreditResponse, + RateLimitResetCreditsSnapshot, + ResetCreditItem, +) +from app.core.config.settings import get_settings +from app.core.crypto import TokenEncryptor +from app.db.models import Account, AccountStatus +from app.db.session import SessionLocal +from app.modules.rate_limit_reset_credits.store import get_rate_limit_reset_credits_store + +pytestmark = pytest.mark.integration + + +@pytest.fixture(autouse=True) +def _reset_settings_cache(): + get_settings.cache_clear() + yield + get_settings.cache_clear() + + +@pytest.fixture(autouse=True) +async def _clear_reset_credit_store(): + await get_rate_limit_reset_credits_store().invalidate() + yield + await get_rate_limit_reset_credits_store().invalidate() + + +def _encode_jwt(payload: dict[str, object]) -> str: + raw = json.dumps(payload, separators=(",", ":")).encode("utf-8") + body = base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + return f"header.{body}.sig" + + +def _make_auth_json(account_id: str, email: str) -> dict[str, object]: + payload: dict[str, object] = { + "email": email, + "chatgpt_account_id": account_id, + "https://api.openai.com/auth": {"chatgpt_plan_type": "plus"}, + } + return { + "tokens": { + "idToken": _encode_jwt(payload), + "accessToken": "access-token", + "refreshToken": "refresh-token", + "accountId": account_id, + }, + } + + +async def _import_account(async_client, account_id: str, email: str) -> str: + auth_json = _make_auth_json(account_id, email) + files = {"auth_json": ("auth.json", json.dumps(auth_json), "application/json")} + response = await async_client.post("/api/accounts/import", files=files) + assert response.status_code == 200 + return generate_unique_account_id(account_id, email) + + +async def _enable_api_key_auth(async_client) -> None: + response = await async_client.put( + "/api/settings", + json={ + "stickyThreadsEnabled": False, + "preferEarlierResetAccounts": False, + "apiKeyAuthEnabled": True, + }, + ) + assert response.status_code == 200 + + +async def _create_api_key(async_client, *, name: str) -> tuple[str, str]: + response = await async_client.post("/api/api-keys/", json={"name": name}) + assert response.status_code == 200 + payload = response.json() + return payload["id"], payload["key"] + + +async def _seed_snapshot( + account_id: str, + *, + available_count: int, + credits: list[ResetCreditItem], +) -> None: + await get_rate_limit_reset_credits_store().set( + account_id, + RateLimitResetCreditsSnapshot( + available_count=available_count, + nearest_expires_at=min( + ( + credit.expires_at + for credit in credits + if credit.status == "available" and credit.expires_at is not None + ), + default=None, + ), + credits=credits, + ), + ) + + +@pytest.mark.asyncio +async def test_v1_reset_credit_requires_valid_bearer_key(async_client): + await _enable_api_key_auth(async_client) + + missing = await async_client.get("/v1/reset-credit") + invalid = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": "Bearer invalid-key"}, + ) + + for response in (missing, invalid): + assert response.status_code == 401 + assert response.json()["error"]["code"] == "invalid_api_key" + + +@pytest.mark.asyncio +async def test_v1_reset_credit_accepts_valid_bearer_key_when_proxy_auth_disabled(async_client): + account_id = await _import_account( + async_client, + "acc-reset-self-service", + "self-service@example.com", + ) + _, key = await _create_api_key(async_client, name="reset-credit-self-service") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-self-service", + status="available", + expires_at=datetime(2031, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + ) + ], + ) + + response = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + ) + + assert response.status_code == 200 + assert response.json() == [ + { + "account_id": account_id, + "email": "self-service@example.com", + "redeem_id": "credit-self-service", + "expiredAt": "2031-01-02T03:04:05Z", + } + ] + + +@pytest.mark.asyncio +async def test_v1_reset_credit_scoped_pool_returns_all_available_credits_for_assigned_account(async_client): + await _enable_api_key_auth(async_client) + assigned_email = "real-assigned@example.com" + other_email = "other@example.com" + assigned_account_id = await _import_account(async_client, "acc-reset-assigned", assigned_email) + other_account_id = await _import_account(async_client, "acc-reset-other", other_email) + + key_id, key = await _create_api_key(async_client, name="reset-credit-scoped") + assign = await async_client.patch( + f"/api/api-keys/{key_id}", + json={"assignedAccountIds": [assigned_account_id]}, + ) + assert assign.status_code == 200 + + soonest = datetime(2031, 1, 2, 3, 4, 5, tzinfo=timezone.utc) + later = soonest + timedelta(hours=2) + await _seed_snapshot( + assigned_account_id, + available_count=2, + credits=[ + ResetCreditItem(id="credit-later", status="available", expires_at=later), + ResetCreditItem(id="credit-soonest", status="available", expires_at=soonest), + ResetCreditItem(id="credit-redeemed", status="redeemed", expires_at=soonest - timedelta(hours=1)), + ], + ) + await _seed_snapshot( + other_account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-other", status="available", expires_at=soonest + timedelta(days=1))], + ) + + response = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + ) + + assert response.status_code == 200 + assert response.json() == [ + { + "account_id": assigned_account_id, + "email": assigned_email, + "redeem_id": "credit-soonest", + "expiredAt": "2031-01-02T03:04:05Z", + }, + { + "account_id": assigned_account_id, + "email": assigned_email, + "redeem_id": "credit-later", + "expiredAt": "2031-01-02T05:04:05Z", + }, + ] + + +@pytest.mark.asyncio +async def test_v1_reset_credit_null_expiry_available_credit_is_returned(async_client): + await _enable_api_key_auth(async_client) + email = "null-expiry@example.com" + account_id = await _import_account(async_client, "acc-reset-null-expiry", email) + + _, key = await _create_api_key(async_client, name="reset-credit-null-expiry") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-null-expiry", status="available", expires_at=None)], + ) + + response = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + ) + + assert response.status_code == 200 + assert response.json() == [ + { + "account_id": account_id, + "email": email, + "redeem_id": "credit-null-expiry", + "expiredAt": None, + } + ] + + +@pytest.mark.asyncio +async def test_v1_reset_credit_mixed_null_expiry_orders_dated_credit_before_null_expiry(async_client): + await _enable_api_key_auth(async_client) + email = "mixed-null-expiry@example.com" + account_id = await _import_account(async_client, "acc-reset-mixed-null-expiry", email) + + _, key = await _create_api_key(async_client, name="reset-credit-mixed-null-expiry") + expires_at = datetime(2031, 2, 1, 1, 2, 3, tzinfo=timezone.utc) + await _seed_snapshot( + account_id, + available_count=2, + credits=[ + ResetCreditItem(id="credit-null-expiry", status="available", expires_at=None), + ResetCreditItem(id="credit-dated", status="available", expires_at=expires_at), + ], + ) + + response = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + ) + + assert response.status_code == 200 + assert response.json() == [ + { + "account_id": account_id, + "email": email, + "redeem_id": "credit-dated", + "expiredAt": "2031-02-01T01:02:03Z", + }, + { + "account_id": account_id, + "email": email, + "redeem_id": "credit-null-expiry", + "expiredAt": None, + }, + ] + + +@pytest.mark.asyncio +async def test_v1_reset_credit_selectable_accounts_excludes_paused_accounts(async_client): + await _enable_api_key_auth(async_client) + active_email = "active@example.com" + paused_email = "paused@example.com" + active_account_id = await _import_account(async_client, "acc-reset-active", active_email) + paused_account_id = await _import_account(async_client, "acc-reset-paused", paused_email) + + pause = await async_client.post( + f"/api/accounts/{paused_account_id}/pause", + json={"reason": "test pause"}, + ) + assert pause.status_code == 200 + + _, key = await _create_api_key(async_client, name="reset-credit-unscoped") + expires_at = datetime(2031, 2, 3, 4, 5, 6, tzinfo=timezone.utc) + await _seed_snapshot( + active_account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-active", status="available", expires_at=expires_at)], + ) + await _seed_snapshot( + paused_account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-paused", status="available", expires_at=expires_at - timedelta(hours=1))], + ) + + response = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + ) + + assert response.status_code == 200 + assert response.json() == [ + { + "account_id": active_account_id, + "email": active_email, + "redeem_id": "credit-active", + "expiredAt": "2031-02-03T04:05:06Z", + } + ] + + +@pytest.mark.asyncio +async def test_v1_reset_credit_excludes_accounts_without_chatgpt_account_id(async_client): + await _enable_api_key_auth(async_client) + active_account_id = await _import_account(async_client, "acc-reset-chatgpt-present", "present@example.com") + missing_id_account_id = await _import_account(async_client, "acc-reset-chatgpt-missing", "missing-id@example.com") + + async with SessionLocal() as session: + await session.execute( + update(Account).where(Account.id == missing_id_account_id).values(chatgpt_account_id=None) + ) + await session.commit() + + _, key = await _create_api_key(async_client, name="reset-credit-chatgpt-filter") + expires_at = datetime(2031, 2, 4, 4, 5, 6, tzinfo=timezone.utc) + await _seed_snapshot( + active_account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-active", status="available", expires_at=expires_at)], + ) + await _seed_snapshot( + missing_id_account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-missing-id", status="available", expires_at=expires_at)], + ) + + response = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + ) + + assert response.status_code == 200 + assert response.json() == [ + { + "account_id": active_account_id, + "email": "present@example.com", + "redeem_id": "credit-active", + "expiredAt": "2031-02-04T04:05:06Z", + } + ] + + +@pytest.mark.asyncio +async def test_v1_reset_credit_duplicate_email_accounts_return_separate_entries(async_client): + await _enable_api_key_auth(async_client) + shared_email = "duplicate@example.com" + first_account_id = await _import_account(async_client, "acc-reset-duplicate-1", shared_email) + second_account_id = await _import_account(async_client, "acc-reset-duplicate-2", shared_email) + + _, key = await _create_api_key(async_client, name="reset-credit-duplicate-email") + first_expires_at = datetime(2031, 3, 4, 5, 6, 7, tzinfo=timezone.utc) + second_expires_at = first_expires_at + timedelta(hours=1) + await _seed_snapshot( + first_account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-duplicate-1", status="available", expires_at=first_expires_at)], + ) + await _seed_snapshot( + second_account_id, + available_count=1, + credits=[ResetCreditItem(id="credit-duplicate-2", status="available", expires_at=second_expires_at)], + ) + + response = await async_client.get( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + ) + + assert response.status_code == 200 + assert response.json() == [ + { + "account_id": first_account_id, + "email": shared_email, + "redeem_id": "credit-duplicate-1", + "expiredAt": "2031-03-04T05:06:07Z", + }, + { + "account_id": second_account_id, + "email": shared_email, + "redeem_id": "credit-duplicate-2", + "expiredAt": "2031-03-04T06:06:07Z", + }, + ] + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_outside_api_key_scope_returns_403(async_client, monkeypatch: pytest.MonkeyPatch): + await _enable_api_key_auth(async_client) + allowed_account_id = await _import_account(async_client, "acc-reset-post-allowed", "allowed@example.com") + blocked_account_id = await _import_account(async_client, "acc-reset-post-blocked", "blocked@example.com") + + key_id, key = await _create_api_key(async_client, name="reset-credit-post-scope") + assign = await async_client.patch( + f"/api/api-keys/{key_id}", + json={"assignedAccountIds": [allowed_account_id]}, + ) + assert assign.status_code == 200 + + consume_mock = AsyncMock() + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", consume_mock) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": blocked_account_id, "redeem_id": "credit-blocked"}, + ) + + assert response.status_code == 403 + assert response.json()["error"]["type"] == "permission_error" + consume_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_rejects_account_without_chatgpt_account_id( + async_client, + monkeypatch: pytest.MonkeyPatch, +): + await _enable_api_key_auth(async_client) + account_id = await _import_account(async_client, "acc-reset-post-missing-chatgpt", "post-missing@example.com") + + async with SessionLocal() as session: + await session.execute(update(Account).where(Account.id == account_id).values(chatgpt_account_id=None)) + await session.commit() + + _, key = await _create_api_key(async_client, name="reset-credit-post-missing-chatgpt") + consume_mock = AsyncMock() + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", consume_mock) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-missing-chatgpt"}, + ) + + assert response.status_code == 403 + assert response.json()["error"]["type"] == "permission_error" + consume_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_unavailable_redeem_id_returns_409(async_client, monkeypatch: pytest.MonkeyPatch): + await _enable_api_key_auth(async_client) + account_id = await _import_account(async_client, "acc-reset-post-missing", "missing@example.com") + + _, key = await _create_api_key(async_client, name="reset-credit-post-missing") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-available", + status="available", + expires_at=datetime(2031, 4, 1, tzinfo=timezone.utc), + ) + ], + ) + + consume_mock = AsyncMock() + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", consume_mock) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-missing"}, + ) + + assert response.status_code == 409 + assert response.json()["error"]["code"] == "invalid_request_error" + consume_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_upstream_conflict_invalidates_stale_snapshot( + async_client, + monkeypatch: pytest.MonkeyPatch, +): + await _enable_api_key_auth(async_client) + account_id = await _import_account(async_client, "acc-reset-post-conflict", "conflict@example.com") + + _, key = await _create_api_key(async_client, name="reset-credit-post-conflict") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-conflict", + status="available", + expires_at=datetime(2031, 4, 2, tzinfo=timezone.utc), + ) + ], + ) + + async def fake_consume(*args, **kwargs): + del args, kwargs + raise ConsumeResetCreditError(409, "credit already redeemed upstream", code="credit_unavailable") + + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", fake_consume) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-conflict"}, + ) + + assert response.status_code == 409 + assert response.json()["error"] == { + "message": "credit already redeemed upstream", + "type": "invalid_request_error", + "code": "invalid_request_error", + } + assert get_rate_limit_reset_credits_store().get(account_id) is None + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_consumes_exact_credit_and_invalidates_snapshot( + async_client, + monkeypatch: pytest.MonkeyPatch, +): + await _enable_api_key_auth(async_client) + email = "exact-credit@example.com" + account_id = await _import_account(async_client, "acc-reset-post-exact", email) + + _, key = await _create_api_key(async_client, name="reset-credit-post-exact") + soonest = datetime(2031, 5, 1, 1, 0, 0, tzinfo=timezone.utc) + later = soonest + timedelta(hours=2) + await _seed_snapshot( + account_id, + available_count=2, + credits=[ + ResetCreditItem(id="credit-soonest", status="available", expires_at=soonest), + ResetCreditItem(id="credit-later", status="available", expires_at=later), + ], + ) + + consume_mock = AsyncMock( + return_value=ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": "credit-later", "status": "redeemed", "redeemed_at": "2031-05-01T03:30:00Z"}, + "windows_reset": 1, + } + ) + ) + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", consume_mock) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-later"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "code": "reset", + "windows_reset": 1, + "redeemed_at": "2031-05-01T03:30:00Z", + } + consume_mock.assert_awaited_once() + consume_args = consume_mock.await_args + assert consume_args is not None + assert consume_args.args[2] == "credit-later" + assert get_rate_limit_reset_credits_store().get(account_id) is None + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_force_refreshes_usage_and_invalidates_selection_cache( + async_client, + monkeypatch: pytest.MonkeyPatch, +): + await _enable_api_key_auth(async_client) + account_id = await _import_account(async_client, "acc-reset-post-refresh", "refresh@example.com") + + _, key = await _create_api_key(async_client, name="reset-credit-post-refresh") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-refresh", + status="available", + expires_at=datetime(2031, 5, 2, 1, 0, 0, tzinfo=timezone.utc), + ) + ], + ) + + consume_mock = AsyncMock( + return_value=ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": "credit-refresh", "status": "redeemed", "redeemed_at": "2031-05-02T01:30:00Z"}, + "windows_reset": 1, + } + ) + ) + force_refresh_calls: list[tuple[str, str]] = [] + + class StubUsageUpdater: + def __init__(self, *args, **kwargs) -> None: + del args, kwargs + + async def force_refresh(self, account) -> bool: + force_refresh_calls.append((account.id, account.status.value)) + return True + + class SelectionCache: + def __init__(self) -> None: + self.invalidations = 0 + + def invalidate(self) -> None: + self.invalidations += 1 + + selection_cache = SelectionCache() + + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", consume_mock) + monkeypatch.setattr("app.modules.proxy.api.UsageUpdater", StubUsageUpdater) + monkeypatch.setattr("app.modules.proxy.api.get_account_selection_cache", lambda: selection_cache) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-refresh"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "code": "reset", + "windows_reset": 1, + "redeemed_at": "2031-05-02T01:30:00Z", + } + assert force_refresh_calls == [(account_id, "active")] + assert selection_cache.invalidations == 1 + assert get_rate_limit_reset_credits_store().get(account_id) is None + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_refreshes_account_before_consuming_credit( + async_client, + monkeypatch: pytest.MonkeyPatch, +): + await _enable_api_key_auth(async_client) + account_id = await _import_account( + async_client, + "acc-reset-post-refresh-token", + "refresh-token@example.com", + ) + + _, key = await _create_api_key(async_client, name="reset-credit-post-refresh-token") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-refresh-token", + status="available", + expires_at=datetime(2031, 5, 2, tzinfo=timezone.utc), + ) + ], + ) + + refreshed_account = SimpleNamespace( + id=account_id, + status=AccountStatus.ACTIVE, + access_token_encrypted=TokenEncryptor().encrypt("fresh-access-token"), + chatgpt_account_id="chatgpt-refresh-token", + ) + events: list[str] = [] + + async def fake_ensure_fresh(requested_account_id: str): + events.append("refresh") + assert requested_account_id == account_id + return refreshed_account + + async def fake_consume( + access_token: str, + chatgpt_account_id: str, + credit_id: str, + *, + route: object | None = None, + allow_direct_egress: bool = False, + ): + events.append("consume") + assert access_token == "fresh-access-token" + assert chatgpt_account_id == "chatgpt-refresh-token" + assert credit_id == "credit-refresh-token" + assert route is None + assert allow_direct_egress is True + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": { + "id": credit_id, + "status": "redeemed", + "redeemed_at": "2031-05-02T00:30:00Z", + }, + "windows_reset": 1, + } + ) + + monkeypatch.setattr("app.modules.proxy.api._ensure_v1_reset_credit_account_fresh", fake_ensure_fresh) + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", fake_consume) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-refresh-token"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "code": "reset", + "windows_reset": 1, + "redeemed_at": "2031-05-02T00:30:00Z", + } + assert events == ["refresh", "consume"] + assert get_rate_limit_reset_credits_store().get(account_id) is None + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_returns_conflict_when_account_refresh_fails( + async_client, + monkeypatch: pytest.MonkeyPatch, +): + await _enable_api_key_auth(async_client) + account_id = await _import_account( + async_client, + "acc-reset-post-refresh-failure", + "refresh-failure@example.com", + ) + + _, key = await _create_api_key(async_client, name="reset-credit-post-refresh-failure") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-refresh-failure", + status="available", + expires_at=datetime(2031, 5, 2, tzinfo=timezone.utc), + ) + ], + ) + + class SelectionCache: + def __init__(self) -> None: + self.invalidations = 0 + + def invalidate(self) -> None: + self.invalidations += 1 + + async def fake_ensure_fresh(requested_account_id: str): + assert requested_account_id == account_id + raise RefreshError("invalid_grant", "refresh token expired", True) + + selection_cache = SelectionCache() + consume_mock = AsyncMock() + + monkeypatch.setattr("app.modules.proxy.api._ensure_v1_reset_credit_account_fresh", fake_ensure_fresh) + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", consume_mock) + monkeypatch.setattr("app.modules.proxy.api.get_account_selection_cache", lambda: selection_cache) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-refresh-failure"}, + ) + + assert response.status_code == 409 + assert response.json() == { + "error": { + "message": "Reset credit redeem could not refresh account credentials: refresh token expired", + "type": "invalid_request_error", + "code": "invalid_request_error", + } + } + assert selection_cache.invalidations == 1 + consume_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_holds_session_open_through_lock_and_upstream_consume( + async_client, + monkeypatch: pytest.MonkeyPatch, +): + await _enable_api_key_auth(async_client) + account_id = await _import_account( + async_client, + "acc-reset-post-session-lifecycle", + "session-lifecycle@example.com", + ) + + _, key = await _create_api_key(async_client, name="reset-credit-post-session-lifecycle") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-session-lifecycle", + status="available", + expires_at=datetime(2031, 6, 1, tzinfo=timezone.utc), + ) + ], + ) + + events: list[str] = [] + repo_sessions = [object(), object()] + account = SimpleNamespace( + id=account_id, + status=AccountStatus.ACTIVE, + access_token_encrypted=TokenEncryptor().encrypt("access-token"), + chatgpt_account_id="chatgpt-session-lifecycle", + ) + + class SessionManager: + async def __aenter__(self): + repo_session = repo_sessions.pop(0) + events.append("session_enter") + return repo_session + + async def __aexit__(self, exc_type, exc, tb): + events.append("session_exit") + return False + + class StubAccountsRepository: + def __init__(self, repo_session_arg): + events.append("repo_init") + + async def get_by_id(self, requested_account_id: str): + events.append("repo_get") + assert requested_account_id == account_id + return account + + @asynccontextmanager + async def fake_serialize_reset_credit_redeem(requested_account_id: str, *, session: object | None): + events.append("lock_wait") + assert requested_account_id == account_id + assert session is not None + events.append("lock_enter") + try: + yield + finally: + events.append("lock_exit") + + async def fake_consume( + access_token: str, + chatgpt_account_id: str, + credit_id: str, + *, + route: object | None = None, + allow_direct_egress: bool = False, + ): + events.append("consume") + assert access_token == "access-token" + assert chatgpt_account_id == "chatgpt-session-lifecycle" + assert credit_id == "credit-session-lifecycle" + assert route is None + assert allow_direct_egress is True + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": { + "id": credit_id, + "status": "redeemed", + "redeemed_at": "2031-06-01T00:30:00Z", + }, + "windows_reset": 1, + } + ) + + async def fake_resolve_route(route_session, requested_account_id: str): + events.append("route_resolve") + assert route_session is not None + assert requested_account_id == account_id + return None + + async def fake_refresh_usage_after_redeem(refreshed_account_id: str) -> None: + events.append("refresh_usage") + assert refreshed_account_id == account_id + + monkeypatch.setattr("app.modules.proxy.api.get_background_session", lambda: SessionManager()) + monkeypatch.setattr("app.modules.proxy.api.AccountsRepository", StubAccountsRepository) + monkeypatch.setattr("app.modules.proxy.api._resolve_reset_credit_route", fake_resolve_route) + monkeypatch.setattr("app.modules.proxy.api.serialize_reset_credit_redeem", fake_serialize_reset_credit_redeem) + monkeypatch.setattr( + "app.modules.proxy.api._ensure_v1_reset_credit_account_fresh", + AsyncMock(side_effect=lambda requested_account_id: events.append("refresh") or account), + ) + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", fake_consume) + monkeypatch.setattr( + "app.modules.proxy.api._refresh_usage_after_v1_reset_credit_redeem", + fake_refresh_usage_after_redeem, + ) + + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-session-lifecycle"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "code": "reset", + "windows_reset": 1, + "redeemed_at": "2031-06-01T00:30:00Z", + } + assert events == [ + "session_enter", + "repo_init", + "repo_get", + "session_exit", + "refresh", + "session_enter", + "repo_init", + "repo_get", + "route_resolve", + "lock_wait", + "lock_enter", + "consume", + "refresh_usage", + "lock_exit", + "session_exit", + ] + assert get_rate_limit_reset_credits_store().get(account_id) is None + + +@pytest.mark.asyncio +async def test_v1_reset_credit_post_preserves_success_when_post_redeem_usage_refresh_fails( + async_client, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + await _enable_api_key_auth(async_client) + account_id = await _import_account( + async_client, + "acc-reset-post-refresh-raise", + "refresh-raise@example.com", + ) + + _, key = await _create_api_key(async_client, name="reset-credit-post-refresh-raise") + await _seed_snapshot( + account_id, + available_count=1, + credits=[ + ResetCreditItem( + id="credit-refresh-raise", + status="available", + expires_at=datetime(2031, 6, 2, tzinfo=timezone.utc), + ) + ], + ) + + account = SimpleNamespace( + id=account_id, + status=AccountStatus.ACTIVE, + access_token_encrypted=TokenEncryptor().encrypt("access-token"), + chatgpt_account_id="chatgpt-refresh-raise", + ) + + async def fake_consume( + access_token: str, + chatgpt_account_id: str, + credit_id: str, + *, + route: object | None = None, + allow_direct_egress: bool = False, + ): + assert access_token == "access-token" + assert chatgpt_account_id == "chatgpt-refresh-raise" + assert credit_id == "credit-refresh-raise" + assert route is None + assert allow_direct_egress is True + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": { + "id": credit_id, + "status": "redeemed", + "redeemed_at": "2031-06-02T00:30:00Z", + }, + "windows_reset": 1, + } + ) + + async def fake_refresh_usage_after_redeem(refreshed_account_id: str) -> None: + assert refreshed_account_id == account_id + raise RuntimeError("usage refresh failed") + + monkeypatch.setattr( + "app.modules.proxy.api._ensure_v1_reset_credit_account_fresh", + AsyncMock(return_value=account), + ) + monkeypatch.setattr("app.modules.proxy.api.consume_reset_credit", fake_consume) + monkeypatch.setattr( + "app.modules.proxy.api._refresh_usage_after_v1_reset_credit_redeem", + fake_refresh_usage_after_redeem, + ) + + with caplog.at_level("WARNING", logger="app.modules.proxy.api"): + response = await async_client.post( + "/v1/reset-credit", + headers={"Authorization": f"Bearer {key}"}, + json={"account_id": account_id, "redeem_id": "credit-refresh-raise"}, + ) + + assert response.status_code == 200 + assert response.json() == { + "code": "reset", + "windows_reset": 1, + "redeemed_at": "2031-06-02T00:30:00Z", + } + assert "V1 reset credit consume succeeded but usage refresh failed" in caplog.text + assert get_rate_limit_reset_credits_store().get(account_id) is None diff --git a/tests/unit/test_rate_limit_reset_credits_api.py b/tests/unit/test_rate_limit_reset_credits_api.py new file mode 100644 index 000000000..2159d3a23 --- /dev/null +++ b/tests/unit/test_rate_limit_reset_credits_api.py @@ -0,0 +1,997 @@ +from __future__ import annotations + +import asyncio +from datetime import datetime +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from fastapi import Request + +from app.core.auth.dependencies import require_dashboard_write_access +from app.core.auth.refresh import RefreshError +from app.core.clients.rate_limit_reset_credits import ( + ConsumeResetCreditError, + ConsumeResetCreditResponse, + RateLimitResetCreditsSnapshot, + ResetCreditFetchError, + ResetCreditItem, + ResetCreditsResponse, +) +from app.core.crypto import TokenEncryptor +from app.core.exceptions import ( + DashboardAuthError, + DashboardConflictError, + DashboardNotFoundError, + DashboardPermissionError, + DashboardServiceUnavailableError, +) +from app.db.models import Account, AccountStatus +from app.modules.rate_limit_reset_credits import api as reset_credits_api +from app.modules.rate_limit_reset_credits.api import ( + ConsumeResetCreditResponseSchema, + _assert_account_can_redeem_reset_credit, + _build_refresh_usage_callback, + _redeem_soonest_reset_credit, + _select_available_credit_by_id, + _select_soonest_available_credit, + _select_soonest_available_credit_from_response, + consume_rate_limit_reset_credit, + get_rate_limit_reset_credits, + serialize_reset_credit_redeem, +) +from app.modules.rate_limit_reset_credits.store import RateLimitResetCreditsStore + +pytestmark = pytest.mark.unit + + +class StubEncryptor(TokenEncryptor): + def __init__(self) -> None: + # Skip key-file I/O; tests only exercise decrypt(). + pass + + def decrypt(self, encrypted: bytes) -> str: + return "decrypted-access-token" + + +def _account(account_id: str = "acc_1") -> Account: + return Account( + id=account_id, + chatgpt_account_id="workspace-1", + email=f"{account_id}@example.com", + plan_type="plus", + access_token_encrypted=b"encrypted", + refresh_token_encrypted=b"refresh", + id_token_encrypted=b"id", + last_refresh=datetime(2025, 1, 1), + status=AccountStatus.ACTIVE, + ) + + +def _credit( + credit_id: str, + *, + status: str = "available", + expires_at: str | None = "2026-07-12T00:00:00Z", +) -> ResetCreditItem: + return ResetCreditItem.model_validate({"id": credit_id, "status": status, "expires_at": expires_at}) + + +def _response(credits: list[ResetCreditItem], available_count: int | None = None) -> ResetCreditsResponse: + count = available_count if available_count is not None else len(credits) + return ResetCreditsResponse(credits=credits, available_count=count) + + +def _fake_request(host: str = "127.0.0.1") -> Request: + return cast(Request, SimpleNamespace(client=SimpleNamespace(host=host))) + + +def _static_fetch_fn(response: ResetCreditsResponse): + async def fetch_fn(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + return response + + return fetch_fn + + +def _snapshot(credits: list[ResetCreditItem], available_count: int | None = None) -> RateLimitResetCreditsSnapshot: + expiries = [ + credit.expires_at for credit in credits if credit.status == "available" and credit.expires_at is not None + ] + return RateLimitResetCreditsSnapshot( + available_count=available_count if available_count is not None else len(credits), + nearest_expires_at=min(expiries) if expiries else None, + credits=credits, + ) + + +# --- GET endpoint --- + + +@pytest.mark.asyncio +async def test_get_returns_null_when_no_snapshot_cached(monkeypatch: pytest.MonkeyPatch) -> None: + store = RateLimitResetCreditsStore() + monkeypatch.setattr(reset_credits_api, "get_rate_limit_reset_credits_store", lambda: store) + + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + return None + + fake_context = SimpleNamespace(repository=_Repo()) + response = await get_rate_limit_reset_credits("acc_missing", context=cast(Any, fake_context)) + assert response is None + + +@pytest.mark.asyncio +async def test_get_populates_cache_on_miss_for_active_account(monkeypatch: pytest.MonkeyPatch) -> None: + store = RateLimitResetCreditsStore() + monkeypatch.setattr(reset_credits_api, "get_rate_limit_reset_credits_store", lambda: store) + + async def _refresh(account, **kwargs: Any) -> None: + await store.set("acc_1", _snapshot([_credit("live")], available_count=1)) + + monkeypatch.setattr(reset_credits_api, "_refresh_account_reset_credits", _refresh) + + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + return _account(account_id) + + fake_context = SimpleNamespace( + repository=_Repo(), + service=SimpleNamespace(_auth_manager=None), + ) + response = await get_rate_limit_reset_credits("acc_1", context=cast(Any, fake_context)) + + assert response is not None + assert response.available_count == 1 + assert response.credits[0].id == "live" + + +@pytest.mark.asyncio +async def test_get_returns_cached_snapshot_shape(monkeypatch: pytest.MonkeyPatch) -> None: + store = RateLimitResetCreditsStore() + await store.set( + "acc_1", + _snapshot([_credit("c1"), _credit("c2", expires_at="2026-06-20T00:00:00Z")], available_count=2), + ) + monkeypatch.setattr(reset_credits_api, "get_rate_limit_reset_credits_store", lambda: store) + + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + return _account(account_id) + + fake_context = SimpleNamespace(repository=_Repo()) + response = await get_rate_limit_reset_credits("acc_1", context=cast(Any, fake_context)) + + assert response is not None + assert response.available_count == 2 + assert response.nearest_expires_at is not None + assert {credit.id for credit in response.credits} == {"c1", "c2"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status", + [AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED], +) +async def test_get_invalidates_cached_snapshot_for_ineligible_status( + status: AccountStatus, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("stale")], available_count=1)) + monkeypatch.setattr(reset_credits_api, "get_rate_limit_reset_credits_store", lambda: store) + + async def _should_not_refresh(*args: Any, **kwargs: Any) -> None: + raise AssertionError("ineligible account should not refresh reset credits") + + monkeypatch.setattr(reset_credits_api, "_refresh_account_reset_credits", _should_not_refresh) + + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + account = _account(account_id) + account.status = status + return account + + fake_context = SimpleNamespace(repository=_Repo()) + + response = await get_rate_limit_reset_credits("acc_1", context=cast(Any, fake_context)) + + assert response is None + assert store.get("acc_1") is None + + +@pytest.mark.asyncio +async def test_get_invalidates_cached_snapshot_without_chatgpt_account_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("stale")], available_count=1)) + monkeypatch.setattr(reset_credits_api, "get_rate_limit_reset_credits_store", lambda: store) + + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + account = _account(account_id) + account.chatgpt_account_id = None + return account + + fake_context = SimpleNamespace(repository=_Repo()) + + response = await get_rate_limit_reset_credits("acc_1", context=cast(Any, fake_context)) + + assert response is None + assert store.get("acc_1") is None + + +# --- soonest-available selection helper --- + + +def test_select_soonest_available_credit_picks_smallest_expires_at() -> None: + credits = [ + _credit("late", expires_at="2026-07-10T00:00:00Z"), + _credit("soon", expires_at="2026-06-20T00:00:00Z"), + _credit("used", status="redeemed", expires_at="2026-06-01T00:00:00Z"), + ] + snapshot = _snapshot(credits) + + selected = _select_soonest_available_credit(snapshot) + + assert selected is not None + assert selected.id == "soon" + + response_selected = _select_soonest_available_credit_from_response(_response(credits)) + assert response_selected is not None + assert response_selected.id == "soon" + + +def test_select_soonest_available_credit_returns_none_when_no_snapshot() -> None: + assert _select_soonest_available_credit(None) is None + + +def test_select_soonest_available_credit_respects_zero_available_count() -> None: + snapshot = _snapshot([_credit("cached_available")], available_count=0) + assert _select_soonest_available_credit(snapshot) is None + + +def test_select_soonest_available_credit_returns_none_when_none_available() -> None: + snapshot = _snapshot([_credit("c1", status="redeemed")]) + assert _select_soonest_available_credit(snapshot) is None + + +def test_select_available_credit_by_id_returns_matching_available_credit() -> None: + response = _response([_credit("wanted"), _credit("other")], available_count=2) + selected = _select_available_credit_by_id(response, "wanted") + assert selected is not None + assert selected.id == "wanted" + + +def test_select_available_credit_by_id_rejects_missing_or_unavailable_credit() -> None: + response = _response([_credit("wanted", status="redeemed")], available_count=1) + assert _select_available_credit_by_id(response, "wanted") is None + assert _select_available_credit_by_id(response, "missing") is None + + +# --- POST consume: helper covers selection, uuid body, invalidation, shape --- + + +@pytest.mark.asyncio +async def test_redeem_returns_409_when_no_available_credit() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("c1", status="redeemed")], available_count=1)) + + with pytest.raises(DashboardConflictError) as excinfo: + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_static_fetch_fn(_response([_credit("c1", status="redeemed")])), + consume_fn=_raise_not_called, # type: ignore[arg-type] + ) + assert excinfo.value.code == "no_available_reset_credit" + cached = store.get("acc_1") + assert cached is not None + assert cached.available_count == 1 + assert cached.credits[0].status == "redeemed" + + +@pytest.mark.asyncio +async def test_redeem_returns_409_when_cached_count_is_zero() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("cached_available")], available_count=0)) + + with pytest.raises(DashboardConflictError) as excinfo: + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_raise_not_called, # type: ignore[arg-type] + consume_fn=_raise_not_called, # type: ignore[arg-type] + ) + assert excinfo.value.code == "no_available_reset_credit" + cached = store.get("acc_1") + assert cached is not None + assert cached.available_count == 0 + assert cached.credits[0].status == "available" + + +@pytest.mark.asyncio +async def test_redeem_returns_409_when_snapshot_missing() -> None: + store = RateLimitResetCreditsStore() + with pytest.raises(DashboardConflictError): + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_raise_not_called, # type: ignore[arg-type] + consume_fn=_raise_not_called, # type: ignore[arg-type] + ) + assert store.get("acc_1") is None + + +@pytest.mark.asyncio +async def test_redeem_replaces_stale_cached_snapshot_when_fresh_fetch_has_no_available_credit() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("stale")], available_count=1)) + + with pytest.raises(DashboardConflictError) as excinfo: + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_static_fetch_fn(_response([], available_count=0)), + consume_fn=_raise_not_called, # type: ignore[arg-type] + ) + + assert excinfo.value.code == "no_available_reset_credit" + cached = store.get("acc_1") + assert cached is not None + assert cached.available_count == 0 + assert cached.credits == [] + + +@pytest.mark.asyncio +async def test_redeem_consumes_fresh_available_credit_when_cached_credit_disappears_upstream() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("stale")], available_count=1)) + + captured: dict[str, Any] = {} + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + captured.update({"access_token": access_token, "account_id": account_id, "credit_id": credit_id}) + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": credit_id, "status": "redeemed", "redeemed_at": "2026-06-13T13:12:31Z"}, + "windows_reset": 1, + } + ) + + result = await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_static_fetch_fn(_response([_credit("other")], available_count=1)), + consume_fn=consume_fn, + ) + + assert captured == { + "access_token": "decrypted-access-token", + "account_id": "workspace-1", + "credit_id": "other", + } + assert result.available_count_before == 1 + assert result.available_count_after == 0 + assert store.get("acc_1") is None + + +@pytest.mark.asyncio +async def test_redeem_reselects_soonest_available_credit_from_fresh_fetch() -> None: + store = RateLimitResetCreditsStore() + await store.set( + "acc_1", + _snapshot([_credit("cached", expires_at="2026-06-30T00:00:00Z")], available_count=1), + ) + + captured: dict[str, Any] = {} + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + captured.update({"access_token": access_token, "account_id": account_id, "credit_id": credit_id}) + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": credit_id, "status": "redeemed", "redeemed_at": "2026-06-13T13:12:31Z"}, + "windows_reset": 1, + } + ) + + result = await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_static_fetch_fn( + _response( + [ + _credit("later", expires_at="2026-07-10T00:00:00Z"), + _credit("fresh-soonest", expires_at="2026-06-20T00:00:00Z"), + ] + ) + ), + consume_fn=consume_fn, + ) + + assert captured == { + "access_token": "decrypted-access-token", + "account_id": "workspace-1", + "credit_id": "fresh-soonest", + } + assert result.available_count_before == 2 + assert result.available_count_after == 1 + assert store.get("acc_1") is None + + +@pytest.mark.asyncio +async def test_redeem_selects_soonest_calls_upstream_and_invalidates_cache() -> None: + store = RateLimitResetCreditsStore() + await store.set( + "acc_1", + _snapshot( + [ + _credit("late", expires_at="2026-07-10T00:00:00Z"), + _credit("soon", expires_at="2026-06-20T00:00:00Z"), + ] + ), + ) + + captured: dict[str, Any] = {} + refreshed: list[str] = [] + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + captured.update({"access_token": access_token, "account_id": account_id, "credit_id": credit_id}) + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": credit_id, "status": "redeemed", "redeemed_at": "2026-06-13T13:12:31Z"}, + "windows_reset": 1, + } + ) + + async def refresh_usage(account: Account) -> None: + refreshed.append(account.id) + + result = await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_static_fetch_fn( + _response( + [ + _credit("late", expires_at="2026-07-10T00:00:00Z"), + _credit("soon", expires_at="2026-06-20T00:00:00Z"), + ] + ) + ), + consume_fn=consume_fn, + refresh_usage=refresh_usage, + ) + + # The soonest-expiring credit id was forwarded with the decrypted token + workspace id. + assert captured == { + "access_token": "decrypted-access-token", + "account_id": "workspace-1", + "credit_id": "soon", + } + # Successful redemption invalidates the in-memory snapshot so the next + # dashboard refresh repulls upstream state instead of serving a local edit. + assert store.get("acc_1") is None + assert result.available_count_before == 2 + assert result.available_count_after == 1 + assert isinstance(result.response, ConsumeResetCreditResponseSchema) + assert result.response.code == "reset" + assert result.response.windows_reset == 1 + assert result.response.redeemed_at is not None + assert result.response.redeemed_at.year == 2026 + assert refreshed == ["acc_1"] + + +@pytest.mark.asyncio +async def test_redeem_restores_snapshot_when_usage_refresh_fails() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("only")], available_count=1)) + fetch_calls = 0 + + async def fetch_fn(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + nonlocal fetch_calls + fetch_calls += 1 + if fetch_calls == 1: + return _response([_credit("only")], available_count=1) + return _response([], available_count=0) + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": credit_id, "status": "redeemed", "redeemed_at": "2026-06-13T13:12:31Z"}, + "windows_reset": 1, + } + ) + + async def refresh_usage(account: Account) -> None: + raise RuntimeError("usage refresh failed") + + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=fetch_fn, + consume_fn=consume_fn, + refresh_usage=refresh_usage, + ) + + restored = store.get("acc_1") + assert fetch_calls == 2 + assert restored is not None + assert restored.available_count == 0 + + +@pytest.mark.asyncio +async def test_redeem_restores_snapshot_when_force_refresh_returns_false( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("only")], available_count=1)) + fetch_calls = 0 + + async def fetch_fn(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + nonlocal fetch_calls + fetch_calls += 1 + if fetch_calls == 1: + return _response([_credit("only")], available_count=1) + return _response([], available_count=0) + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": credit_id, "status": "redeemed", "redeemed_at": "2026-06-13T13:12:31Z"}, + "windows_reset": 1, + } + ) + + class _UsageUpdater: + async def force_refresh(self, account: Account) -> bool: + assert account.id == "acc_1" + return False + + class _SelectionCache: + def __init__(self) -> None: + self.invalidated = 0 + + def invalidate(self) -> None: + self.invalidated += 1 + + selection_cache = _SelectionCache() + monkeypatch.setattr(reset_credits_api, "get_account_selection_cache", lambda: selection_cache) + refresh_usage = _build_refresh_usage_callback( + cast(Any, SimpleNamespace(service=SimpleNamespace(_usage_updater=_UsageUpdater()))) + ) + + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=fetch_fn, + consume_fn=consume_fn, + refresh_usage=refresh_usage, + ) + + restored = store.get("acc_1") + assert fetch_calls == 2 + assert restored is not None + assert restored.available_count == 0 + assert selection_cache.invalidated == 0 + + +@pytest.mark.asyncio +async def test_redeem_serializes_requests_per_account() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("only")], available_count=1)) + fetch_calls = 0 + + async def fetch_fn(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + nonlocal fetch_calls + fetch_calls += 1 + if fetch_calls == 1: + return _response([_credit("only")], available_count=1) + return _response([], available_count=0) + + started = asyncio.Event() + release = asyncio.Event() + consume_calls: list[str] = [] + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + consume_calls.append(credit_id) + started.set() + await release.wait() + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": credit_id, "status": "redeemed", "redeemed_at": "2026-06-13T13:12:31Z"}, + "windows_reset": 1, + } + ) + + first = asyncio.create_task( + _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=fetch_fn, + consume_fn=consume_fn, + ) + ) + await started.wait() + + second = asyncio.create_task( + _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=fetch_fn, + consume_fn=consume_fn, + ) + ) + await asyncio.sleep(0) + + assert consume_calls == ["only"] + + release.set() + await first + + with pytest.raises(DashboardConflictError) as excinfo: + await second + assert excinfo.value.code == "no_available_reset_credit" + assert consume_calls == ["only"] + + +@pytest.mark.asyncio +async def test_serialize_reset_credit_redeem_uses_postgresql_advisory_lock() -> None: + calls: list[tuple[str, dict[str, str] | None]] = [] + + class _FakeSession: + def get_bind(self) -> Any: + return SimpleNamespace(dialect=SimpleNamespace(name="postgresql")) + + async def execute(self, statement: Any, params: dict[str, str] | None = None) -> None: + calls.append((str(statement), params)) + + async with serialize_reset_credit_redeem("acc_1", session=cast(Any, _FakeSession())): + pass + + assert calls == [ + ( + "SELECT pg_advisory_xact_lock(hashtext(:lock_key))", + {"lock_key": "reset-credit-redeem:acc_1"}, + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "expected_exception"), + [ + (401, DashboardAuthError), + (403, DashboardPermissionError), + (409, DashboardConflictError), + (503, DashboardServiceUnavailableError), + (0, DashboardServiceUnavailableError), + ], +) +async def test_redeem_translates_upstream_fetch_failures( + status_code: int, + expected_exception: type[Exception], +) -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("only")], available_count=1)) + + async def fetch_fn(*args: Any, **kwargs: Any) -> ResetCreditsResponse: + raise ResetCreditFetchError(status_code, f"upstream fetch failed {status_code}", code=f"fetch_{status_code}") + + with pytest.raises(expected_exception) as excinfo: + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=fetch_fn, + consume_fn=_raise_not_called, # type: ignore[arg-type] + ) + + assert str(excinfo.value) == f"upstream fetch failed {status_code}" + assert getattr(excinfo.value, "code", None) == f"fetch_{status_code}" + cached = store.get("acc_1") + assert cached is not None + assert [credit.id for credit in cached.credits] == ["only"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "expected_exception"), + [ + (401, DashboardAuthError), + (403, DashboardPermissionError), + (409, DashboardConflictError), + (503, DashboardServiceUnavailableError), + (0, DashboardServiceUnavailableError), + ], +) +async def test_redeem_translates_upstream_consume_failures( + status_code: int, + expected_exception: type[Exception], +) -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("only")], available_count=1)) + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + raise ConsumeResetCreditError(status_code, f"upstream failed {status_code}", code=f"upstream_{status_code}") + + with pytest.raises(expected_exception) as excinfo: + await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_static_fetch_fn(_response([_credit("only")], available_count=1)), + consume_fn=consume_fn, + ) + + assert str(excinfo.value) == f"upstream failed {status_code}" + assert getattr(excinfo.value, "code", None) == f"upstream_{status_code}" + cached = store.get("acc_1") + assert cached is not None + assert [credit.id for credit in cached.credits] == ["only"] + + +@pytest.mark.asyncio +async def test_redeem_reports_zero_available_count_after_last_credit() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_1", _snapshot([_credit("only")], available_count=1)) + + async def consume_fn( + access_token: str, + account_id: str | None, + credit_id: str, + **kwargs: Any, + ) -> ConsumeResetCreditResponse: + return ConsumeResetCreditResponse.model_validate( + { + "code": "reset", + "credit": {"id": credit_id, "status": "redeemed", "redeemed_at": "2026-06-13T13:12:31Z"}, + "windows_reset": 1, + } + ) + + result = await _redeem_soonest_reset_credit( + account=_account(), + store=store, + encryptor=StubEncryptor(), + fetch_fn=_static_fetch_fn(_response([_credit("only")], available_count=1)), + consume_fn=consume_fn, + ) + + assert result.available_count_after == 0 + + +@pytest.mark.parametrize( + "status", + [AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED], +) +def test_assert_account_can_redeem_reset_credit_rejects_non_applicable_statuses(status: AccountStatus) -> None: + account = _account() + account.status = status + with pytest.raises(DashboardConflictError) as excinfo: + _assert_account_can_redeem_reset_credit(account) + assert excinfo.value.code == "account_not_reset_credit_applicable" + + +def test_assert_account_can_redeem_reset_credit_rejects_missing_chatgpt_account_id() -> None: + account = _account() + account.chatgpt_account_id = None + with pytest.raises(DashboardConflictError) as excinfo: + _assert_account_can_redeem_reset_credit(account) + assert excinfo.value.code == "account_not_reset_credit_applicable" + + +# --- POST consume: handler-level 404 when account missing --- + + +@pytest.mark.asyncio +async def test_consume_handler_returns_404_when_account_missing() -> None: + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + return None + + fake_context = SimpleNamespace(repository=_Repo()) + + with pytest.raises(DashboardNotFoundError): + await consume_rate_limit_reset_credit( + _fake_request(), + account_id="missing", + _write_access=None, + context=cast(Any, fake_context), + ) + + +@pytest.mark.asyncio +async def test_consume_handler_audits_live_available_count_before_when_cache_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = RateLimitResetCreditsStore() + monkeypatch.setattr(reset_credits_api, "get_rate_limit_reset_credits_store", lambda: store) + + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + return _account(account_id) + + logged: dict[str, Any] = {} + + def _log_async(event: str, **kwargs: Any) -> None: + logged["event"] = event + logged.update(kwargs) + + async def _redeem(**kwargs: Any) -> Any: + return reset_credits_api._RedeemResetCreditOutcome( + response=ConsumeResetCreditResponseSchema(code="reset", windows_reset=1, redeemed_at=None), + available_count_before=3, + available_count_after=2, + ) + + monkeypatch.setattr(reset_credits_api, "_redeem_soonest_reset_credit", _redeem) + monkeypatch.setattr(reset_credits_api.AuditService, "log_async", _log_async) + + fake_context = SimpleNamespace( + repository=_Repo(), + service=SimpleNamespace(_auth_manager=None, _usage_updater=None), + ) + response = await consume_rate_limit_reset_credit( + _fake_request(), + account_id="acc_1", + _write_access=None, + context=cast(Any, fake_context), + ) + + assert response.code == "reset" + assert logged["event"] == "account_rate_limit_reset_credit_consumed" + assert logged["details"]["available_reset_credits_before"] == 3 + assert logged["details"]["available_reset_credits_after"] == 2 + + +@pytest.mark.asyncio +async def test_consume_handler_invalidates_selection_cache_on_permanent_refresh_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + return _account(account_id) + + class _SelectionCache: + def __init__(self) -> None: + self.invalidated = 0 + + def invalidate(self) -> None: + self.invalidated += 1 + + async def _redeem(**kwargs: Any) -> Any: + raise RefreshError("invalid_grant", "refresh token expired", True) + + selection_cache = _SelectionCache() + monkeypatch.setattr(reset_credits_api, "_redeem_soonest_reset_credit", _redeem) + monkeypatch.setattr(reset_credits_api, "get_account_selection_cache", lambda: selection_cache) + + fake_context = SimpleNamespace( + repository=_Repo(), + service=SimpleNamespace(_auth_manager=None, _usage_updater=None), + ) + with pytest.raises(DashboardConflictError) as excinfo: + await consume_rate_limit_reset_credit( + _fake_request(), + account_id="acc_1", + _write_access=None, + context=cast(Any, fake_context), + ) + + assert excinfo.value.code == "account_reset_credit_refresh_failed" + assert selection_cache.invalidated == 1 + + +@pytest.mark.asyncio +async def test_consume_handler_keeps_selection_cache_on_transient_refresh_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _Repo: + async def get_by_id(self, account_id: str) -> Account | None: + return _account(account_id) + + class _SelectionCache: + def __init__(self) -> None: + self.invalidated = 0 + + def invalidate(self) -> None: + self.invalidated += 1 + + async def _redeem(**kwargs: Any) -> Any: + raise RefreshError("transport_error", "timeout", False) + + selection_cache = _SelectionCache() + monkeypatch.setattr(reset_credits_api, "_redeem_soonest_reset_credit", _redeem) + monkeypatch.setattr(reset_credits_api, "get_account_selection_cache", lambda: selection_cache) + + fake_context = SimpleNamespace( + repository=_Repo(), + service=SimpleNamespace(_auth_manager=None, _usage_updater=None), + ) + with pytest.raises(DashboardConflictError) as excinfo: + await consume_rate_limit_reset_credit( + _fake_request(), + account_id="acc_1", + _write_access=None, + context=cast(Any, fake_context), + ) + + assert excinfo.value.code == "account_reset_credit_refresh_failed" + assert selection_cache.invalidated == 0 + + +# --- POST consume: write-access gating refuses guests (full ASGI path) --- + + +@pytest.mark.asyncio +async def test_consume_refuses_read_only_guest(app_instance, async_client) -> None: # type: ignore[no-untyped-def] + async def _guest_refused(_request: Any = None) -> None: + raise DashboardPermissionError( + "Read-only dashboard access cannot modify dashboard state", + code="read_only_access", + ) + + app_instance.dependency_overrides[require_dashboard_write_access] = _guest_refused + try: + response = await async_client.post("/api/accounts/acc_guest/rate-limit-reset-credits/consume") + finally: + app_instance.dependency_overrides.pop(require_dashboard_write_access, None) + + assert response.status_code == 403 + + +async def _raise_not_called(*args: Any, **kwargs: Any) -> Any: + raise AssertionError("consume_fn must not be called when no credit is available") diff --git a/tests/unit/test_rate_limit_reset_credits_client.py b/tests/unit/test_rate_limit_reset_credits_client.py new file mode 100644 index 000000000..68f82709e --- /dev/null +++ b/tests/unit/test_rate_limit_reset_credits_client.py @@ -0,0 +1,480 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, cast +from uuid import UUID + +import pytest + +from app.core.clients.headers import build_chatgpt_auth_headers +from app.core.clients.rate_limit_reset_credits import ( + ConsumeResetCreditError, + ConsumeResetCreditResponse, + ResetCreditFetchError, + ResetCreditsResponse, + build_snapshot, + consume_reset_credit, + fetch_reset_credits, +) +from app.core.clients.usage import _usage_headers +from app.core.config.settings import get_settings + +pytestmark = pytest.mark.unit + + +class StubResponse: + def __init__(self, status: int, payload: object | None, text: str) -> None: + self.status = status + self._payload = payload + self._text = text + + async def json(self, content_type: str | None = None) -> object: + if self._payload is None: + raise ValueError("no json") + return self._payload + + async def text(self) -> str: + return self._text + + +@dataclass +class ClientState: + calls: int = 0 + method: str | None = None + url: str | None = None + headers: dict[str, str] | None = None + json_body: dict[str, Any] | None = None + + +class StubRequestContext: + def __init__( + self, + responses: list[StubResponse], + state: ClientState, + method: str, + url: str, + headers: dict[str, str], + json_body: dict[str, Any] | None, + retry_options: object | None, + ) -> None: + self._responses = responses + self._state = state + self._method = method + self._url = url + self._headers = headers + self._json_body = json_body + self._retry_options = retry_options + + async def __aenter__(self) -> StubResponse: + attempts = getattr(self._retry_options, "attempts", 1) + statuses = set(getattr(self._retry_options, "statuses", set())) + response: StubResponse | None = None + for attempt in range(attempts): + index = min(self._state.calls, len(self._responses) - 1) + response = self._responses[index] + self._state.calls += 1 + self._state.method = self._method + self._state.url = self._url + self._state.headers = dict(self._headers) + self._state.json_body = dict(self._json_body) if self._json_body else None + if response.status in statuses and attempt < attempts - 1: + continue + return response + if response is None: + response = StubResponse(500, None, "no response") + return response + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + +class StubRetryClient: + def __init__(self, responses: list[StubResponse], state: ClientState) -> None: + self._responses = responses + self._state = state + + def request( + self, + method: str, + url: str, + headers: dict[str, str] | None = None, + json: dict[str, Any] | None = None, + timeout: object | None = None, + retry_options: object | None = None, + ) -> StubRequestContext: + return StubRequestContext( + self._responses, + self._state, + method, + url, + headers or {}, + json, + retry_options, + ) + + +def _list_payload() -> dict: + return { + "credits": [ + { + "id": "RateLimitResetCredit_test", + "reset_type": "codex_rate_limits", + "status": "available", + "granted_at": "2026-06-12T01:29:41.346025Z", + "expires_at": "2026-07-12T01:29:41.346025Z", + "redeem_started_at": None, + "redeemed_at": None, + "title": "One free rate limit reset", + "description": "Thanks for using Codex!", + } + ], + "available_count": 1, + } + + +def test_usage_headers_delegate_to_shared_helper() -> None: + """The usage client and the reset-credits client share one header builder.""" + assert _usage_headers("tok", "acc_workspace") == build_chatgpt_auth_headers("tok", "acc_workspace") + + +@pytest.mark.asyncio +async def test_fetch_reset_credits_sends_bearer_and_account_id_headers() -> None: + state = ClientState() + client = StubRetryClient([StubResponse(200, _list_payload(), "")], state) + + data = await fetch_reset_credits( + "access-token", + "acc_workspace", + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert isinstance(data, ResetCreditsResponse) + assert data.available_count == 1 + assert data.credits[0].id == "RateLimitResetCredit_test" + assert state.method == "GET" + assert state.url == "http://upstream.test/backend-api/wham/rate-limit-reset-credits" + assert state.headers is not None + assert state.headers["Authorization"] == "Bearer access-token" + assert state.headers["chatgpt-account-id"] == "acc_workspace" + + +@pytest.mark.asyncio +async def test_fetch_reset_credits_skips_account_id_header_for_email_and_local_prefixes() -> None: + for account_id in ("email_user@example.com", "local_abcd"): + state = ClientState() + client = StubRetryClient([StubResponse(200, _list_payload(), "")], state) + await fetch_reset_credits( + "access-token", + account_id, + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + assert state.headers is not None + assert "chatgpt-account-id" not in state.headers, account_id + + +@pytest.mark.asyncio +async def test_fetch_reset_credits_normalizes_base_url_without_backend_api_segment() -> None: + state = ClientState() + client = StubRetryClient([StubResponse(200, {"credits": [], "available_count": 0}, "")], state) + + await fetch_reset_credits( + "access-token", + None, + base_url="http://upstream.test/", # trailing slash, no backend-api + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert state.url == "http://upstream.test/backend-api/wham/rate-limit-reset-credits" + + +@pytest.mark.asyncio +async def test_fetch_reset_credits_raises_on_non_200() -> None: + state = ClientState() + client = StubRetryClient( + [StubResponse(401, {"error": {"code": "unauthorized", "message": "bad token"}}, "")], + state, + ) + + with pytest.raises(ResetCreditFetchError) as excinfo: + await fetch_reset_credits( + "access-token", + None, + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 401 + assert excinfo.value.code == "unauthorized" + + +@pytest.mark.asyncio +async def test_fetch_reset_credits_handles_non_json_body() -> None: + state = ClientState() + client = StubRetryClient([StubResponse(502, None, "boom")], state) + + with pytest.raises(ResetCreditFetchError) as excinfo: + await fetch_reset_credits( + "access-token", + None, + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_fetch_reset_credits_rejects_malformed_success_body() -> None: + state = ClientState() + client = StubRetryClient([StubResponse(200, ["not", "an", "object"], "")], state) + + with pytest.raises(ResetCreditFetchError) as excinfo: + await fetch_reset_credits( + "access-token", + None, + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_fetch_reset_credits_rejects_success_body_missing_contract_fields() -> None: + state = ClientState() + client = StubRetryClient([StubResponse(200, {"credits": []}, "")], state) + + with pytest.raises(ResetCreditFetchError) as excinfo: + await fetch_reset_credits( + "access-token", + None, + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_consume_reset_credit_sends_credit_id_and_uuid_redeem_request_id() -> None: + state = ClientState() + client = StubRetryClient( + [ + StubResponse( + 200, + { + "code": "reset", + "credit": { + "id": "RateLimitResetCredit_test", + "status": "redeemed", + "redeemed_at": "2026-06-13T13:12:31Z", + }, + "windows_reset": 1, + }, + "", + ) + ], + state, + ) + + result = await consume_reset_credit( + "access-token", + "acc_workspace", + "RateLimitResetCredit_test", + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert isinstance(result, ConsumeResetCreditResponse) + assert result.code == "reset" + assert result.windows_reset == 1 + assert result.credit is not None and result.credit.redeemed_at is not None + assert state.method == "POST" + assert state.url == "http://upstream.test/backend-api/wham/rate-limit-reset-credits/consume" + assert state.headers is not None + assert state.headers["Authorization"] == "Bearer access-token" + assert state.headers["chatgpt-account-id"] == "acc_workspace" + assert state.headers["Content-Type"] == "application/json" + # body carries the credit id and a freshly-generated uuid redeem_request_id + assert state.json_body is not None + assert state.json_body["credit_id"] == "RateLimitResetCredit_test" + redeem_request_id = state.json_body["redeem_request_id"] + assert isinstance(redeem_request_id, str) and len(redeem_request_id) == 36 + # canonical uuid v4 + parsed = UUID(redeem_request_id, version=4) + assert str(parsed) == redeem_request_id + + +@pytest.mark.asyncio +async def test_consume_reset_credit_generates_fresh_redeem_request_id_each_call() -> None: + ids: list[str] = [] + for _ in range(2): + state = ClientState() + client = StubRetryClient( + [StubResponse(200, {"code": "reset", "credit": {"id": "x"}, "windows_reset": 1}, "")], + state, + ) + await consume_reset_credit( + "access-token", + None, + "RateLimitResetCredit_test", + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + assert state.json_body is not None + ids.append(state.json_body["redeem_request_id"]) + assert ids[0] != ids[1] + + +@pytest.mark.asyncio +async def test_consume_reset_credit_does_not_retry_when_max_retries_omitted(monkeypatch: pytest.MonkeyPatch) -> None: + settings = get_settings() + original_retries = settings.usage_fetch_max_retries + monkeypatch.setattr(settings, "usage_fetch_max_retries", 2) + + state = ClientState() + client = StubRetryClient( + [ + StubResponse(503, {"error": {"code": "temporarily_unavailable", "message": "retry later"}}, ""), + StubResponse(200, {"code": "reset", "credit": {"id": "x"}, "windows_reset": 1}, ""), + ], + state, + ) + + with pytest.raises(ConsumeResetCreditError) as excinfo: + await consume_reset_credit( + "access-token", + None, + "RateLimitResetCredit_test", + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 503 + assert excinfo.value.code == "temporarily_unavailable" + assert state.calls == 1 + monkeypatch.setattr(settings, "usage_fetch_max_retries", original_retries) + + +@pytest.mark.asyncio +async def test_consume_reset_credit_raises_on_non_200() -> None: + state = ClientState() + client = StubRetryClient( + [StubResponse(409, {"error": {"code": "no_credit", "message": "none"}}, "")], + state, + ) + + with pytest.raises(ConsumeResetCreditError) as excinfo: + await consume_reset_credit( + "access-token", + None, + "RateLimitResetCredit_test", + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 409 + assert excinfo.value.code == "no_credit" + + +@pytest.mark.asyncio +async def test_consume_reset_credit_rejects_malformed_success_body() -> None: + state = ClientState() + client = StubRetryClient([StubResponse(200, "not json", "")], state) + + with pytest.raises(ConsumeResetCreditError) as excinfo: + await consume_reset_credit( + "access-token", + None, + "RateLimitResetCredit_test", + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_consume_reset_credit_rejects_success_body_missing_contract_fields() -> None: + state = ClientState() + client = StubRetryClient([StubResponse(200, {"code": "reset"}, "")], state) + + with pytest.raises(ConsumeResetCreditError) as excinfo: + await consume_reset_credit( + "access-token", + None, + "RateLimitResetCredit_test", + base_url="http://upstream.test/backend-api", + timeout_seconds=2.0, + max_retries=0, + client=cast(Any, client), + allow_direct_egress=True, + ) + + assert excinfo.value.status_code == 502 + + +def test_build_snapshot_projects_nearest_available_expiry() -> None: + response = ResetCreditsResponse.model_validate( + { + "credits": [ + {"id": "a", "status": "available", "expires_at": "2026-07-10T00:00:00Z"}, + {"id": "b", "status": "available", "expires_at": "2026-06-20T00:00:00Z"}, + {"id": "c", "status": "redeemed", "expires_at": "2026-06-01T00:00:00Z"}, + ], + "available_count": 2, + } + ) + + snapshot = build_snapshot(response) + + assert snapshot.available_count == 2 + assert snapshot.nearest_expires_at is not None + assert snapshot.nearest_expires_at.year == 2026 + assert snapshot.nearest_expires_at.month == 6 + assert snapshot.nearest_expires_at.day == 20 + assert [credit.id for credit in snapshot.credits] == ["a", "b", "c"] + + +def test_build_snapshot_returns_none_expiry_when_no_available_credit() -> None: + response = ResetCreditsResponse.model_validate( + {"credits": [{"id": "a", "status": "redeemed"}], "available_count": 0} + ) + assert build_snapshot(response).nearest_expires_at is None diff --git a/tests/unit/test_rate_limit_reset_credits_mapper.py b/tests/unit/test_rate_limit_reset_credits_mapper.py new file mode 100644 index 000000000..8503ceafb --- /dev/null +++ b/tests/unit/test_rate_limit_reset_credits_mapper.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from datetime import datetime +from typing import cast + +import pytest + +from app.core.clients.rate_limit_reset_credits import RateLimitResetCreditsSnapshot +from app.core.crypto import TokenEncryptor +from app.db.models import Account, AccountStatus +from app.modules.accounts.mappers import build_account_summaries +from app.modules.rate_limit_reset_credits.store import RateLimitResetCreditsStore + +pytestmark = pytest.mark.unit + +_DEFAULT_CHATGPT_ACCOUNT_ID = object() + + +def _account( + account_id: str, + *, + status: AccountStatus = AccountStatus.ACTIVE, + chatgpt_account_id: str | None | object = _DEFAULT_CHATGPT_ACCOUNT_ID, +) -> Account: + if chatgpt_account_id is _DEFAULT_CHATGPT_ACCOUNT_ID: + resolved_chatgpt_account_id: str | None = f"workspace-{account_id}" + else: + resolved_chatgpt_account_id = cast("str | None", chatgpt_account_id) + return Account( + id=account_id, + chatgpt_account_id=resolved_chatgpt_account_id, + email=f"{account_id}@example.com", + plan_type="plus", + access_token_encrypted=b"", + refresh_token_encrypted=b"", + id_token_encrypted=b"", + last_refresh=datetime(2025, 1, 1), + status=status, + ) + + +def _summaries(accounts: list[Account], store: RateLimitResetCreditsStore): + return build_account_summaries( + accounts=accounts, + primary_usage={}, + secondary_usage={}, + encryptor=TokenEncryptor(), + include_auth=False, + reset_credits_store=store, + ) + + +def test_account_summary_exposes_cached_reset_credits_fields() -> None: + store = RateLimitResetCreditsStore() + nearest = datetime(2026, 7, 10, 0, 0, 0) + store_snapshot = RateLimitResetCreditsSnapshot( + available_count=2, + nearest_expires_at=nearest, + credits=[], + ) + # Bypass the async lock by writing the backing dict directly for a unit fixture. + store._snapshots["acc_with_credits"] = store_snapshot # type: ignore[attr-defined] + + [summary] = _summaries([_account("acc_with_credits")], store) + + assert summary.available_reset_credits == 2 + assert summary.reset_credit_nearest_expires_at == nearest + + +def test_account_summary_returns_zero_and_null_when_no_snapshot() -> None: + store = RateLimitResetCreditsStore() + + [summary] = _summaries([_account("acc_no_cache")], store) + + assert summary.available_reset_credits == 0 + assert summary.reset_credit_nearest_expires_at is None + + +@pytest.mark.parametrize( + "status", + [AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED], +) +def test_account_summary_suppresses_cached_reset_credits_for_ineligible_status(status: AccountStatus) -> None: + store = RateLimitResetCreditsStore() + store._snapshots["acc_ineligible"] = RateLimitResetCreditsSnapshot( # type: ignore[attr-defined] + available_count=3, + nearest_expires_at=datetime(2026, 6, 20, 0, 0, 0), + credits=[], + ) + + [summary] = _summaries([_account("acc_ineligible", status=status)], store) + + assert summary.available_reset_credits == 0 + assert summary.reset_credit_nearest_expires_at is None + + +def test_account_summary_suppresses_cached_reset_credits_without_chatgpt_account_id() -> None: + store = RateLimitResetCreditsStore() + store._snapshots["acc_no_workspace"] = RateLimitResetCreditsSnapshot( # type: ignore[attr-defined] + available_count=3, + nearest_expires_at=datetime(2026, 6, 20, 0, 0, 0), + credits=[], + ) + + [summary] = _summaries([_account("acc_no_workspace", chatgpt_account_id=None)], store) + + assert summary.available_reset_credits == 0 + assert summary.reset_credit_nearest_expires_at is None + + +def test_account_summary_mixed_cache_state_across_accounts() -> None: + store = RateLimitResetCreditsStore() + store._snapshots["acc_has"] = RateLimitResetCreditsSnapshot( # type: ignore[attr-defined] + available_count=5, + nearest_expires_at=datetime(2026, 6, 20, 0, 0, 0), + credits=[], + ) + + summaries = _summaries([_account("acc_has"), _account("acc_missing")], store) + by_id = {s.account_id: s for s in summaries} + + assert by_id["acc_has"].available_reset_credits == 5 + assert by_id["acc_has"].reset_credit_nearest_expires_at is not None + assert by_id["acc_missing"].available_reset_credits == 0 + assert by_id["acc_missing"].reset_credit_nearest_expires_at is None + + +def test_account_summary_does_not_crash_when_store_is_empty() -> None: + store = RateLimitResetCreditsStore() + accounts = [_account(f"acc_{i}") for i in range(3)] + + summaries = _summaries(accounts, store) + + assert len(summaries) == 3 + assert all(s.available_reset_credits == 0 for s in summaries) + assert all(s.reset_credit_nearest_expires_at is None for s in summaries) diff --git a/tests/unit/test_rate_limit_reset_credits_scheduler.py b/tests/unit/test_rate_limit_reset_credits_scheduler.py new file mode 100644 index 000000000..a331c82a4 --- /dev/null +++ b/tests/unit/test_rate_limit_reset_credits_scheduler.py @@ -0,0 +1,387 @@ +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from datetime import datetime +from typing import Any + +import pytest + +from app.core.clients.rate_limit_reset_credits import ( + RateLimitResetCreditsSnapshot, + ResetCreditFetchError, + ResetCreditsResponse, +) +from app.core.crypto import TokenEncryptor +from app.core.usage import reset_credits_refresh_scheduler as scheduler_module +from app.core.usage.reset_credits_refresh_scheduler import ( + RateLimitResetCreditsRefreshScheduler, + refresh_reset_credits_for_accounts, +) +from app.db.models import Account, AccountStatus +from app.modules.rate_limit_reset_credits.store import RateLimitResetCreditsStore + +pytestmark = pytest.mark.unit + + +class StubEncryptor(TokenEncryptor): + def __init__(self) -> None: + # Skip key-file I/O; tests only exercise decrypt(). + pass + + def decrypt(self, encrypted: bytes) -> str: + return f"token-for-{encrypted.decode() if encrypted else ''}" + + +def _make_account( + account_id: str, + *, + status: AccountStatus = AccountStatus.ACTIVE, + chatgpt_account_id: str | None = "workspace-x", +) -> Account: + return Account( + id=account_id, + chatgpt_account_id=chatgpt_account_id, + email=f"{account_id}@example.com", + plan_type="plus", + access_token_encrypted=account_id.encode(), + refresh_token_encrypted=b"refresh", + id_token_encrypted=b"id", + last_refresh=datetime(2025, 1, 1), + status=status, + ) + + +def _response(available_count: int = 1) -> ResetCreditsResponse: + return ResetCreditsResponse.model_validate( + { + "credits": [ + {"id": "c1", "status": "available", "expires_at": "2026-07-12T00:00:00Z"}, + ], + "available_count": available_count, + } + ) + + +@pytest.mark.asyncio +async def test_refresh_skips_paused_reauth_and_deactivated_accounts() -> None: + store = RateLimitResetCreditsStore() + stale = RateLimitResetCreditsSnapshot(available_count=5) + await store.set("acc_paused", stale) + await store.set("acc_reauth", stale) + await store.set("acc_deactivated", stale) + fetched: list[str] = [] + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + fetched.append(access_token) + return _response() + + accounts = [ + _make_account("acc_paused", status=AccountStatus.PAUSED), + _make_account("acc_reauth", status=AccountStatus.REAUTH_REQUIRED), + _make_account("acc_deactivated", status=AccountStatus.DEACTIVATED), + _make_account("acc_active"), + ] + + await refresh_reset_credits_for_accounts( + accounts=accounts, + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + + # Only the active account was fetched and cached. + assert fetched == ["token-for-acc_active"] + assert store.get("acc_paused") is stale + assert store.get("acc_reauth") is stale + assert store.get("acc_deactivated") is stale + assert store.get("acc_active") is not None + + +@pytest.mark.asyncio +async def test_refresh_skips_account_without_chatgpt_account_id() -> None: + store = RateLimitResetCreditsStore() + stale = RateLimitResetCreditsSnapshot(available_count=4) + await store.set("acc_no_workspace", stale) + fetched: list[str] = [] + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + fetched.append(access_token) + return _response() + + await refresh_reset_credits_for_accounts( + accounts=[_make_account("acc_no_workspace", chatgpt_account_id=None)], + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + + assert fetched == [] + assert store.get("acc_no_workspace") is stale + + +@pytest.mark.asyncio +async def test_refresh_401_retains_prior_snapshot_without_status_mutation() -> None: + """A 401 on reset-credits must not trigger a token refresh or status write. + + Reset-credits polling owns no account-status derivation; usage refresh owns + token refresh and deactivation. A 401 logs and retains the prior cached + snapshot with a single fetch attempt and no AuthManager involvement. + """ + store = RateLimitResetCreditsStore() + prior = RateLimitResetCreditsSnapshot(available_count=2) + await store.set("acc_401", prior) + account = _make_account("acc_401", status=AccountStatus.ACTIVE) + fetch_calls = {"count": 0} + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + fetch_calls["count"] += 1 + raise ResetCreditFetchError(401, "unauthorized") + + await refresh_reset_credits_for_accounts( + accounts=[account], + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + + assert fetch_calls["count"] == 1 + assert store.get("acc_401") is prior + assert account.status == AccountStatus.ACTIVE + + +@pytest.mark.asyncio +async def test_one_account_failure_does_not_break_the_loop() -> None: + store = RateLimitResetCreditsStore() + fetched: list[str] = [] + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + fetched.append(access_token) + if access_token == "token-for-acc_fail": + raise ResetCreditFetchError(500, "boom") + return _response(available_count=3) + + accounts = [_make_account("acc_fail"), _make_account("acc_ok")] + + await refresh_reset_credits_for_accounts( + accounts=accounts, + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + + # Both accounts were attempted despite the first raising. + assert fetched == ["token-for-acc_fail", "token-for-acc_ok"] + # The failing account left no snapshot; the healthy one was cached. + assert store.get("acc_fail") is None + ok_snapshot = store.get("acc_ok") + assert ok_snapshot is not None + assert ok_snapshot.available_count == 3 + + +@pytest.mark.asyncio +async def test_upstream_error_retains_prior_snapshot_and_does_not_mutate_status() -> None: + store = RateLimitResetCreditsStore() + prior = RateLimitResetCreditsSnapshot(available_count=2) + await store.set("acc_retain", prior) + account = _make_account("acc_retain", status=AccountStatus.ACTIVE) + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + raise ResetCreditFetchError(503, "busy") + + await refresh_reset_credits_for_accounts( + accounts=[account], + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + + # Prior snapshot is retained exactly. + assert store.get("acc_retain") is prior + assert prior.available_count == 2 + # Account status is untouched. + assert account.status == AccountStatus.ACTIVE + + +@pytest.mark.asyncio +async def test_refresh_does_not_resurrect_snapshot_invalidated_during_fetch() -> None: + store = RateLimitResetCreditsStore() + prior = RateLimitResetCreditsSnapshot(available_count=1) + await store.set("acc_redeemed", prior) + account = _make_account("acc_redeemed", status=AccountStatus.ACTIVE) + fetch_started = asyncio.Event() + release_fetch = asyncio.Event() + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + fetch_started.set() + await release_fetch.wait() + return _response(available_count=1) + + refresh_task = asyncio.create_task( + refresh_reset_credits_for_accounts( + accounts=[account], + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + ) + await fetch_started.wait() + + await store.invalidate("acc_redeemed") + release_fetch.set() + await refresh_task + + assert store.get("acc_redeemed") is None + + +@pytest.mark.asyncio +async def test_unrelated_account_write_does_not_drop_in_flight_refresh() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_a", RateLimitResetCreditsSnapshot(available_count=1)) + account = _make_account("acc_b", status=AccountStatus.ACTIVE) + fetch_started = asyncio.Event() + release_fetch = asyncio.Event() + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + fetch_started.set() + await release_fetch.wait() + return _response(available_count=4) + + refresh_task = asyncio.create_task( + refresh_reset_credits_for_accounts( + accounts=[account], + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + ) + await fetch_started.wait() + + await store.set("acc_a", RateLimitResetCreditsSnapshot(available_count=9)) + release_fetch.set() + await refresh_task + + snapshot_b = store.get("acc_b") + assert snapshot_b is not None + assert snapshot_b.available_count == 4 + + +@pytest.mark.asyncio +async def test_refresh_never_calls_account_status_writes() -> None: + """The scheduler must not transition account status under any path. + + The refresh function operates only on the in-memory store; it holds no + reference to a repository and therefore cannot perform status writes. We + assert the account objects are byte-identical in status before and after, + including across the failure path. + """ + store = RateLimitResetCreditsStore() + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + if access_token == "token-for-acc_fail": + raise ResetCreditFetchError(401, "unauthorized") + return _response() + + accounts = [_make_account("acc_fail"), _make_account("acc_ok")] + statuses_before = {a.id: a.status for a in accounts} + + await refresh_reset_credits_for_accounts( + accounts=accounts, + encryptor=StubEncryptor(), + store=store, + fetch_fn=fetch_fn, + ) + + assert {a.id: a.status for a in accounts} == statuses_before + + +@pytest.mark.asyncio +async def test_refresh_once_caches_snapshots_on_each_replica(monkeypatch: pytest.MonkeyPatch) -> None: + """Each process refreshes its own in-memory cache without leader gating.""" + + account = _make_account("acc_replica") + store = RateLimitResetCreditsStore() + + captured: list[Any] = [] + + class _FakeRepo: + async def list_accounts(self) -> list[Account]: + captured.append("list_accounts") + return [account] + + class _FakeSession: + async def __aenter__(self) -> "_FakeSession": + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + return None + + @asynccontextmanager + async def _fake_background_session(): + captured.append("session_opened") + yield _FakeSession() + + monkeypatch.setattr(scheduler_module, "get_background_session", _fake_background_session) + monkeypatch.setattr(scheduler_module, "AccountsRepository", lambda session: _FakeRepo()) + monkeypatch.setattr(scheduler_module, "TokenEncryptor", lambda: StubEncryptor()) + monkeypatch.setattr(scheduler_module, "get_rate_limit_reset_credits_store", lambda: store) + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + captured.append(("fetch", access_token, account_id)) + return _response(available_count=7) + + monkeypatch.setattr(scheduler_module, "fetch_reset_credits", fetch_fn) + + scheduler = RateLimitResetCreditsRefreshScheduler(interval_seconds=60) + await scheduler._refresh_once() + + assert ("fetch", "token-for-acc_replica", "workspace-x") in captured + snapshot = store.get("acc_replica") + assert snapshot is not None + assert snapshot.available_count == 7 + assert account.status == AccountStatus.ACTIVE + + +@pytest.mark.asyncio +async def test_refresh_once_releases_list_session_before_account_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + account = _make_account("acc_release") + store = RateLimitResetCreditsStore() + events: list[str] = [] + + class _FakeRepo: + async def list_accounts(self) -> list[Account]: + events.append("list_accounts") + return [account] + + class _FakeSession: + async def __aenter__(self) -> "_FakeSession": + events.append("session_enter") + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + events.append("session_exit") + return None + + @asynccontextmanager + async def _fake_background_session(): + events.append("session_enter") + try: + yield _FakeSession() + finally: + events.append("session_exit") + + monkeypatch.setattr(scheduler_module, "get_background_session", _fake_background_session) + monkeypatch.setattr(scheduler_module, "AccountsRepository", lambda session: _FakeRepo()) + monkeypatch.setattr(scheduler_module, "TokenEncryptor", lambda: StubEncryptor()) + monkeypatch.setattr(scheduler_module, "get_rate_limit_reset_credits_store", lambda: store) + + async def fetch_fn(access_token: str, account_id: str | None, **kwargs: Any) -> ResetCreditsResponse: + events.append("fetch") + return _response() + + monkeypatch.setattr(scheduler_module, "fetch_reset_credits", fetch_fn) + + scheduler = RateLimitResetCreditsRefreshScheduler(interval_seconds=60) + await scheduler._refresh_once() + + assert events == ["session_enter", "list_accounts", "session_exit", "fetch"] diff --git a/tests/unit/test_rate_limit_reset_credits_store.py b/tests/unit/test_rate_limit_reset_credits_store.py new file mode 100644 index 000000000..51504691b --- /dev/null +++ b/tests/unit/test_rate_limit_reset_credits_store.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +from datetime import datetime + +import pytest + +from app.core.clients.rate_limit_reset_credits import RateLimitResetCreditsSnapshot, ResetCreditItem +from app.modules.rate_limit_reset_credits.store import ( + RateLimitResetCreditsStore, + get_rate_limit_reset_credits_store, +) + +pytestmark = pytest.mark.unit + + +def _snapshot(available_count: int = 1) -> RateLimitResetCreditsSnapshot: + return RateLimitResetCreditsSnapshot(available_count=available_count) + + +def _credit(credit_id: str, *, expires_at: str, status: str = "available") -> ResetCreditItem: + return ResetCreditItem.model_validate({"id": credit_id, "expires_at": expires_at, "status": status}) + + +@pytest.mark.asyncio +async def test_set_and_get_round_trip() -> None: + store = RateLimitResetCreditsStore() + snapshot = _snapshot(2) + + await store.set("acc_a", snapshot) + + assert store.get("acc_a") is snapshot + assert snapshot.available_count == 2 + + +@pytest.mark.asyncio +async def test_get_returns_none_for_missing_account() -> None: + store = RateLimitResetCreditsStore() + assert store.get("missing") is None + + +@pytest.mark.asyncio +async def test_set_overwrites_prior_snapshot() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_a", _snapshot(1)) + await store.set("acc_a", _snapshot(5)) + + snapshot = store.get("acc_a") + assert snapshot is not None + assert snapshot.available_count == 5 + + +@pytest.mark.asyncio +async def test_generation_changes_are_scoped_to_account() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_a", _snapshot(1)) + await store.set("acc_b", _snapshot(2)) + generation_b = store.generation("acc_b") + + await store.set("acc_a", _snapshot(9)) + + assert await store.set_if_generation("acc_b", _snapshot(7), generation_b) + snapshot_b = store.get("acc_b") + assert snapshot_b is not None + assert snapshot_b.available_count == 7 + + +@pytest.mark.asyncio +async def test_same_account_generation_change_rejects_stale_write() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_a", _snapshot(1)) + generation = store.generation("acc_a") + + await store.invalidate("acc_a") + + assert not await store.set_if_generation("acc_a", _snapshot(7), generation) + assert store.get("acc_a") is None + + +@pytest.mark.asyncio +async def test_invalidate_all_rejects_in_flight_writes_for_any_account() -> None: + store = RateLimitResetCreditsStore() + generation = store.generation("acc_a") + + await store.invalidate() + + assert not await store.set_if_generation("acc_a", _snapshot(7), generation) + assert store.get("acc_a") is None + + +@pytest.mark.asyncio +async def test_invalidate_single_account_clears_only_that_key() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_a", _snapshot(1)) + await store.set("acc_b", _snapshot(2)) + + await store.invalidate("acc_a") + + assert store.get("acc_a") is None + snapshot_b = store.get("acc_b") + assert snapshot_b is not None + assert snapshot_b.available_count == 2 + + +@pytest.mark.asyncio +async def test_invalidate_all_clears_every_key() -> None: + store = RateLimitResetCreditsStore() + await store.set("acc_a", _snapshot(1)) + await store.set("acc_b", _snapshot(2)) + + await store.invalidate() + + assert store.get("acc_a") is None + assert store.get("acc_b") is None + + +@pytest.mark.asyncio +async def test_invalidate_missing_account_is_noop() -> None: + store = RateLimitResetCreditsStore() + await store.invalidate("never_existed") # must not raise + assert store.get("never_existed") is None + + +@pytest.mark.asyncio +async def test_mark_credit_redeemed_preserves_remaining_available_credits() -> None: + store = RateLimitResetCreditsStore() + await store.set( + "acc_a", + RateLimitResetCreditsSnapshot( + available_count=2, + nearest_expires_at=datetime.fromisoformat("2026-06-20T00:00:00+00:00"), + credits=[ + _credit("soon", expires_at="2026-06-20T00:00:00Z"), + _credit("late", expires_at="2026-07-10T00:00:00Z"), + ], + ), + ) + redeemed_at = datetime.fromisoformat("2026-06-18T12:00:00+00:00") + + await store.mark_credit_redeemed("acc_a", "soon", redeemed_at=redeemed_at) + + snapshot = store.get("acc_a") + assert snapshot is not None + assert snapshot.available_count == 1 + assert snapshot.nearest_expires_at == datetime.fromisoformat("2026-07-10T00:00:00+00:00") + assert [(credit.id, credit.status) for credit in snapshot.credits] == [("soon", "redeemed"), ("late", "available")] + assert snapshot.credits[0].redeemed_at == redeemed_at + + +@pytest.mark.asyncio +async def test_concurrent_setters_are_serialized_under_lock() -> None: + store = RateLimitResetCreditsStore() + + async def writer(account_id: str) -> None: + for value in range(20): + await store.set(account_id, _snapshot(value)) + + # If the lock did not serialize, a careless implementation could still pass, + # but a dict is not coroutine-safe across truly concurrent writes; this at + # least exercises the lock path and confirms the final state is consistent. + import asyncio + + await asyncio.gather(*(writer(f"acc_{i}") for i in range(5))) + + for i in range(5): + snapshot = store.get(f"acc_{i}") + assert snapshot is not None + assert snapshot.available_count == 19 + + +def test_module_singleton_accessor_returns_shared_instance() -> None: + assert get_rate_limit_reset_credits_store() is get_rate_limit_reset_credits_store()