src/providers/chatgpt/provider.py

@@ -16,7 +16,7 @@ from .tokens import (
     load_tokens,
     promote_next_tokens,
     refresh_tokens,
-    save_next_tokens,
+    save_state,
     save_tokens,
 )
 from .usage import get_usage_data
@@ -44,10 +44,14 @@ class ChatGPTProvider(Provider):
                 attempt,
                 CHATGPT_REGISTRATION_MAX_ATTEMPTS,
             )
-            success = await self.register_new_account()
-            if success:
+            generated_tokens = await register_chatgpt_account(
+                email_provider_factory=self.email_provider_factory,
+            )
+            if generated_tokens:
+                save_tokens(generated_tokens)
                 return True
             logger.warning("Registration attempt %s failed", attempt)
+            await asyncio.sleep(1.5 * attempt)
         return False
 
     async def _create_next_account_under_lock(self) -> bool:
@@ -56,23 +60,28 @@ class ChatGPTProvider(Provider):
             return True
 
         logger.info("Creating next account")
-        success = await self._register_with_retries()
-        if not success:
-            return False
+        for attempt in range(1, CHATGPT_REGISTRATION_MAX_ATTEMPTS + 1):
+            logger.info(
+                "Next-account registration attempt %s/%s",
+                attempt,
+                CHATGPT_REGISTRATION_MAX_ATTEMPTS,
+            )
+            generated_tokens = await register_chatgpt_account(
+                email_provider_factory=self.email_provider_factory,
+            )
+            if generated_tokens:
+                if active_before:
+                    save_state(active_before, generated_tokens)
+                else:
+                    save_state(generated_tokens, None)
+                logger.info("Next account is ready")
+                return True
+            logger.warning("Next-account registration attempt %s failed", attempt)
+            await asyncio.sleep(1.5 * attempt)
 
-        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
+        if active_before or next_before:
+            save_state(active_before, next_before)
+        return False
 
     async def force_recreate_token(self) -> str | None:
         async with self._token_write_lock:
@@ -85,6 +94,9 @@ class ChatGPTProvider(Provider):
                 return None
             return tokens.access_token
 
+    async def startup_prepare(self) -> None:
+        await self.ensure_next_account()
+
     async def ensure_next_account(self) -> bool:
         next_tokens = load_next_tokens()
         if next_tokens and not next_tokens.is_expired:
@@ -96,6 +108,14 @@ class ChatGPTProvider(Provider):
                 return True
             return await self._create_next_account_under_lock()
 
+    async def ensure_standby_account(
+        self,
+        usage_percent: int,
+        prepare_threshold: int,
+    ) -> None:
+        if usage_percent >= prepare_threshold:
+            await self.ensure_next_account()
+
     async def maybe_switch_active_account(self, usage_percent: int) -> bool:
         if usage_percent < CHATGPT_SWITCH_THRESHOLD:
             return False
@@ -119,6 +139,9 @@ class ChatGPTProvider(Provider):
                 )
             return switched
 
+    async def maybe_rotate_account(self, usage_percent: int) -> bool:
+        return await self.maybe_switch_active_account(usage_percent)
+
     @property
     def name(self) -> str:
         return "chatgpt"
@@ -154,13 +177,17 @@ class ChatGPTProvider(Provider):
 
     async def register_new_account(self) -> bool:
         """Register a new ChatGPT account"""
-        return await register_chatgpt_account(
+        generated_tokens = await register_chatgpt_account(
             email_provider_factory=self.email_provider_factory,
         )
+        if not generated_tokens:
+            return False
+        save_tokens(generated_tokens)
+        return True
 
     async def get_usage_info(self, access_token: str) -> dict[str, Any]:
         """Get usage information for the current token"""
-        usage_data = get_usage_data(access_token)
+        usage_data = await get_usage_data(access_token)
         if not usage_data:
             return {"error": "Failed to get usage"}