From 0ca2a44885a99c7b765d444d3bb9b6011e44c359 Mon Sep 17 00:00:00 2001
From: TheSmallHanCat <w190909332@outlook.com>
Date: Tue, 2 Dec 2025 23:39:26 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E9=94=99=E8=AF=AF=E8=AE=A1=E6=95=B0?=
 =?UTF-8?q?=E7=BB=9F=E8=AE=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/core/database.py          | 25 +++++++++++++++++++++----
 src/core/models.py            |  3 ++-
 src/services/token_manager.py |  6 +++---
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/src/core/database.py b/src/core/database.py
index c172d74..8778e01 100644
--- a/src/core/database.py
+++ b/src/core/database.py
@@ -204,6 +204,20 @@ class Database:
                         except Exception as e:
                             print(f"  ✗ Failed to add column '{col_name}': {e}")
 
+            # Check and add missing columns to token_stats table
+            if await self._table_exists(db, "token_stats"):
+                columns_to_add = [
+                    ("consecutive_error_count", "INTEGER DEFAULT 0"),
+                ]
+
+                for col_name, col_type in columns_to_add:
+                    if not await self._column_exists(db, "token_stats", col_name):
+                        try:
+                            await db.execute(f"ALTER TABLE token_stats ADD COLUMN {col_name} {col_type}")
+                            print(f"  ✓ Added column '{col_name}' to token_stats table")
+                        except Exception as e:
+                            print(f"  ✗ Failed to add column '{col_name}': {e}")
+
             # Check and add missing columns to admin_config table
             if await self._table_exists(db, "admin_config"):
                 columns_to_add = [
@@ -291,6 +305,7 @@ class Database:
                     today_video_count INTEGER DEFAULT 0,
                     today_error_count INTEGER DEFAULT 0,
                     today_date DATE,
+                    consecutive_error_count INTEGER DEFAULT 0,
                     FOREIGN KEY (token_id) REFERENCES tokens(id)
                 )
             """)
@@ -825,7 +840,7 @@ class Database:
             await db.commit()
     
     async def increment_error_count(self, token_id: int):
-        """Increment error count"""
+        """Increment error count (both total and consecutive)"""
         from datetime import date
         async with aiosqlite.connect(self.db_path) as db:
             today = str(date.today())
@@ -838,16 +853,18 @@ class Database:
                 await db.execute("""
                     UPDATE token_stats
                     SET error_count = error_count + 1,
+                        consecutive_error_count = consecutive_error_count + 1,
                         today_error_count = 1,
                         today_date = ?,
                         last_error_at = CURRENT_TIMESTAMP
                     WHERE token_id = ?
                 """, (today, token_id))
             else:
-                # Same day, just increment both
+                # Same day, just increment all counters
                 await db.execute("""
                     UPDATE token_stats
                     SET error_count = error_count + 1,
+                        consecutive_error_count = consecutive_error_count + 1,
                         today_error_count = today_error_count + 1,
                         today_date = ?,
                         last_error_at = CURRENT_TIMESTAMP
@@ -856,10 +873,10 @@ class Database:
             await db.commit()
     
     async def reset_error_count(self, token_id: int):
-        """Reset error count"""
+        """Reset consecutive error count (keep total error_count)"""
         async with aiosqlite.connect(self.db_path) as db:
             await db.execute("""
-                UPDATE token_stats SET error_count = 0 WHERE token_id = ?
+                UPDATE token_stats SET consecutive_error_count = 0 WHERE token_id = ?
             """, (token_id,))
             await db.commit()
     
diff --git a/src/core/models.py b/src/core/models.py
index d47e59a..806e09e 100644
--- a/src/core/models.py
+++ b/src/core/models.py
@@ -43,12 +43,13 @@ class TokenStats(BaseModel):
     token_id: int
     image_count: int = 0
     video_count: int = 0
-    error_count: int = 0
+    error_count: int = 0  # Historical total errors (never reset)
     last_error_at: Optional[datetime] = None
     today_image_count: int = 0
     today_video_count: int = 0
     today_error_count: int = 0
     today_date: Optional[str] = None
+    consecutive_error_count: int = 0  # Consecutive errors for auto-disable
 
 class Task(BaseModel):
     """Task model"""
diff --git a/src/services/token_manager.py b/src/services/token_manager.py
index c5b5cf9..be88cd5 100644
--- a/src/services/token_manager.py
+++ b/src/services/token_manager.py
@@ -913,12 +913,12 @@ class TokenManager:
     async def record_error(self, token_id: int):
         """Record token error"""
         await self.db.increment_error_count(token_id)
-        
+
         # Check if should ban
         stats = await self.db.get_token_stats(token_id)
         admin_config = await self.db.get_admin_config()
-        
-        if stats and stats.error_count >= admin_config.error_ban_threshold:
+
+        if stats and stats.consecutive_error_count >= admin_config.error_ban_threshold:
             await self.db.update_token_status(token_id, False)
     
     async def record_success(self, token_id: int, is_video: bool = False):