refactor: add staged next-account rotation and clarify ChatGPT config

src/providers/chatgpt/provider.py

@@ -1,19 +1,30 @@
 import asyncio
 import logging
-from typing import Callable
+import os
 from typing import Any
+from typing import Callable
 
 from playwright.async_api import BrowserContext
 
-from providers.base import Provider, ProviderTokens
 from email_providers import BaseProvider
 from email_providers import TempMailOrgProvider
-from .tokens import load_tokens, save_tokens, refresh_tokens
+from providers.base import Provider, ProviderTokens
+from .tokens import (
+    clear_next_tokens,
+    load_next_tokens,
+    load_state,
+    load_tokens,
+    promote_next_tokens,
+    refresh_tokens,
+    save_next_tokens,
+    save_tokens,
+)
 from .usage import get_usage_data
 from .registration import register_chatgpt_account
 
 logger = logging.getLogger(__name__)
-MAX_REGISTRATION_ATTEMPTS = 4
+CHATGPT_REGISTRATION_MAX_ATTEMPTS = 4
+CHATGPT_SWITCH_THRESHOLD = int(os.environ.get("CHATGPT_SWITCH_THRESHOLD", "95"))
 
 
 class ChatGPTProvider(Provider):
@@ -27,11 +38,11 @@ class ChatGPTProvider(Provider):
         self._token_write_lock = asyncio.Lock()
 
     async def _register_with_retries(self) -> bool:
-        for attempt in range(1, MAX_REGISTRATION_ATTEMPTS + 1):
+        for attempt in range(1, CHATGPT_REGISTRATION_MAX_ATTEMPTS + 1):
             logger.info(
                 "Registration attempt %s/%s",
                 attempt,
-                MAX_REGISTRATION_ATTEMPTS,
+                CHATGPT_REGISTRATION_MAX_ATTEMPTS,
             )
             success = await self.register_new_account()
             if success:
@@ -39,16 +50,75 @@ class ChatGPTProvider(Provider):
             logger.warning("Registration attempt %s failed", attempt)
         return False
 
+    async def _create_next_account_under_lock(self) -> bool:
+        active_before, next_before = load_state()
+        if next_before:
+            return True
+
+        logger.info("Creating next account")
+        success = await self._register_with_retries()
+        if not success:
+            return False
+
+        generated_active = load_tokens()
+        if not generated_active:
+            return False
+
+        # Registration writes new tokens as active; restore old active and keep
+        # generated account as next.
+        if active_before:
+            save_tokens(active_before)
+        else:
+            clear_next_tokens()
+        save_next_tokens(generated_active)
+        logger.info("Next account is ready")
+        return True
+
     async def force_recreate_token(self) -> str | None:
         async with self._token_write_lock:
             success = await self._register_with_retries()
             if not success:
                 return None
+            clear_next_tokens()
             tokens = load_tokens()
             if not tokens:
                 return None
             return tokens.access_token
 
+    async def ensure_next_account(self) -> bool:
+        next_tokens = load_next_tokens()
+        if next_tokens and not next_tokens.is_expired:
+            return True
+
+        async with self._token_write_lock:
+            next_tokens = load_next_tokens()
+            if next_tokens and not next_tokens.is_expired:
+                return True
+            return await self._create_next_account_under_lock()
+
+    async def maybe_switch_active_account(self, usage_percent: int) -> bool:
+        if usage_percent < CHATGPT_SWITCH_THRESHOLD:
+            return False
+
+        async with self._token_write_lock:
+            next_tokens = load_next_tokens()
+            if not next_tokens or next_tokens.is_expired:
+                logger.info(
+                    "Active usage >= %s%% and next account missing",
+                    CHATGPT_SWITCH_THRESHOLD,
+                )
+                created = await self._create_next_account_under_lock()
+                if not created:
+                    return False
+
+            switched = promote_next_tokens()
+            if switched:
+                logger.info(
+                    "Switched active account (usage >= %s%%)",
+                    CHATGPT_SWITCH_THRESHOLD,
+                )
+            return switched
+
     @property
     def name(self) -> str:
         return "chatgpt"