From a0c294ab45fb4aa7889299bbe8b33812096e0f62 Mon Sep 17 00:00:00 2001
From: Holegots <fuergaosi@gmail.com>
Date: Fri, 9 Dec 2022 19:01:04 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20account=20relogin=20if=20refr?=
 =?UTF-8?q?esh=20token=20failed?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md      |  4 +--
 README_ZH.md   |  4 +--
 src/cache.ts   |  6 ++++-
 src/chatgpt.ts | 69 +++++++++++++++++++++++++++++++++++++++++++++++---
 4 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 3460afe..2f422b9 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,7 @@
 </p>
 
 > Use ChatGPT On Wechat via wechaty
-
+English | [中文文档](README.zh-CN.md)
 ### 🏠 [Homepage](https://github.com/fuergaosi233/wechat-chatgpt)
 
 ## 🌟 Feature
@@ -22,7 +22,7 @@
 - [x] Add Dockerfile
 - [x] Publish to Docker.hub
 - [ ] Add Railway deploy
-- [ ] Auto Reload OpenAI Accounts Pool
+- [x] Auto Reload OpenAI Accounts Pool
 - [ ] Add sendmessage retry for 429/503
 
 ## Use with docker (recommended)
diff --git a/README_ZH.md b/README_ZH.md
index 98b5648..61c71c8 100644
--- a/README_ZH.md
+++ b/README_ZH.md
@@ -10,7 +10,7 @@
 </p>
 
 > 在微信上迅速接入 ChatGPT，让它成为你最好的助手！
-
+[English](README.md) | 中文文档
 ### 🏠 [主页](https://github.com/fuergaosi233/wechat-chatgpt/blob/main/README_ZH.md)
 
 ## 🌟 功能点
@@ -22,7 +22,7 @@
 - [x] 加入 Dockerfile
 - [x] 发布到 Docker.hub
 - [ ] 通过 Railway 进行部署
-- [ ] 实现 OpenAI 账户池的热加载
+- [x] 实现 OpenAI 账户池的热加载
 - [ ] 当 OpenAI 返回码为 429/503 时自动重试
 
 ## 通过 Docker 使用（✅ 推荐）
diff --git a/src/cache.ts b/src/cache.ts
index bbe3dbf..18440f3 100644
--- a/src/cache.ts
+++ b/src/cache.ts
@@ -12,8 +12,12 @@ export class Cache {
   get(key: string) {
     return this._cache[key];
   }
-  set(key: string, value: any) {
+  async set(key: string, value: any) {
     this._cache[key] = value;
     fs.writeFileSync(this._file, JSON.stringify(this._cache));
   }
+  async delete(key: string) {
+    delete this._cache[key];
+    fs.writeFileSync(this._file, JSON.stringify(this._cache));
+  }
 }
diff --git a/src/chatgpt.ts b/src/chatgpt.ts
index 71281c8..3d7876d 100644
--- a/src/chatgpt.ts
+++ b/src/chatgpt.ts
@@ -10,6 +10,8 @@ import {
   AccountWithUserInfo,
   isAccountWithUserInfo,
   isAccountWithSessionToken,
+  AccountWithSessionToken,
+  IAccount,
 } from "./interface.js";
 
 enum MessageType {
@@ -53,7 +55,6 @@ export class ChatGPTPoole {
     }
     const cmd = `poetry run python3 src/generate_session.py ${email} ${password}`;
     const platform = process.platform;
-
     const { stdout, stderr, exitCode } = await execa(
       platform === "win32" ? "powershell" : "sh",
       [platform === "win32" ? "/c" : "-c", cmd]
@@ -70,6 +71,64 @@ export class ChatGPTPoole {
     }
     return "";
   }
+  async resetAccount(account: IAccount) {
+    if (isAccountWithUserInfo(account)) {
+      // Remove all conversation information
+      this.conversationsPool.forEach((item, key) => {
+        if ((item.account as AccountWithUserInfo)?.email === account.email) {
+          this.conversationsPool.delete(key);
+        }
+      });
+      // Relogin and generate a new session token
+      const chatGPTItem = this.chatGPTPools
+        .filter((item) => isAccountWithUserInfo(item.account))
+        .find(
+          (
+            item: any
+          ): item is IChatGPTItem & {
+            account: AccountWithUserInfo;
+            chatGpt: ChatGPTAPI;
+          } => item?.email === account.email
+        );
+      if (chatGPTItem) {
+        await this.cache.delete(account.email);
+        try {
+          const session_token = await this.getSessionToken(
+            chatGPTItem.account?.email,
+            chatGPTItem.account?.password
+          );
+          chatGPTItem.chatGpt = new ChatGPTAPI({
+            sessionToken: session_token,
+          });
+        } catch (err) {
+          //remove this object
+          this.chatGPTPools = this.chatGPTPools.filter(
+            (item) =>
+              (item.account as AccountWithUserInfo)?.email !== account.email
+          );
+          console.error(
+            `Try reset account: ${account.email} failed: ${err}, remove it from pool`
+          );
+        }
+      }
+    } else if (isAccountWithSessionToken(account)) {
+      // Remove all conversation information
+      this.conversationsPool.forEach((item, key) => {
+        if (
+          (item.account as AccountWithSessionToken)?.session_token ===
+          account.session_token
+        ) {
+          this.conversationsPool.delete(key);
+        }
+      });
+      // Remove this gptItem
+      this.chatGPTPools = this.chatGPTPools.filter(
+        (item) =>
+          (item.account as AccountWithSessionToken)?.session_token !==
+          account.session_token
+      );
+    }
+  }
   async startPools() {
     const sessionAccounts = config.chatGPTAccountPool.filter(
       isAccountWithSessionToken
@@ -125,7 +184,7 @@ export class ChatGPTPoole {
     return conversationItem;
   }
   // send message with talkid
-  async sendMessage(message: string, talkid: string) {
+  async sendMessage(message: string, talkid: string): Promise<string> {
     const conversationItem = this.getConversation(talkid);
     const { conversation, account } = conversationItem;
     try {
@@ -133,6 +192,10 @@ export class ChatGPTPoole {
       const response = await conversation.sendMessage(message);
       return response;
     } catch (err: any) {
+      if (err.message.includes("ChatGPT failed to refresh auth token")) {
+        await this.resetAccount(account);
+        return this.sendMessage(message, talkid);
+      }
       console.error(
         `err is ${err.message}, account ${JSON.stringify(account)}`
       );
@@ -166,7 +229,7 @@ export class ChatGPTBot {
     return `@${this.botName}`;
   }
   async startGPTBot() {
-    console.debug(`Start GPT Bot Config is:${config}`);
+    console.debug(`Start GPT Bot Config is:${JSON.stringify(config)}`);
     await this.chatGPTPool.startPools();
     console.debug(`🤖️ Start GPT Bot Success, ready to handle message!`);
   }