mirror of
https://github.com/TheSmallHanCat/sora2api.git
synced 2026-02-13 17:34:42 +08:00
feat: 增强Token禁用状态管理,区分失效与禁用
This commit is contained in:
@@ -249,7 +249,10 @@ async def get_tokens(token: str = Depends(verify_admin_token)) -> List[dict]:
|
||||
"video_enabled": token.video_enabled,
|
||||
# 并发限制
|
||||
"image_concurrency": token.image_concurrency,
|
||||
"video_concurrency": token.video_concurrency
|
||||
"video_concurrency": token.video_concurrency,
|
||||
# 过期和禁用信息
|
||||
"is_expired": token.is_expired,
|
||||
"disabled_reason": token.disabled_reason
|
||||
})
|
||||
|
||||
return result
|
||||
|
||||
@@ -393,7 +393,8 @@ class Database:
|
||||
video_enabled BOOLEAN DEFAULT 1,
|
||||
image_concurrency INTEGER DEFAULT -1,
|
||||
video_concurrency INTEGER DEFAULT -1,
|
||||
is_expired BOOLEAN DEFAULT 0
|
||||
is_expired BOOLEAN DEFAULT 0,
|
||||
disabled_reason TEXT
|
||||
)
|
||||
""")
|
||||
|
||||
@@ -586,6 +587,22 @@ class Database:
|
||||
if not await self._column_exists(db, "admin_config", "auto_disable_on_401"):
|
||||
await db.execute("ALTER TABLE admin_config ADD COLUMN auto_disable_on_401 BOOLEAN DEFAULT 1")
|
||||
|
||||
# Migration: Add disabled_reason column to tokens table if it doesn't exist
|
||||
if not await self._column_exists(db, "tokens", "disabled_reason"):
|
||||
await db.execute("ALTER TABLE tokens ADD COLUMN disabled_reason TEXT")
|
||||
# For existing disabled tokens without a reason, set to 'manual'
|
||||
await db.execute("""
|
||||
UPDATE tokens
|
||||
SET disabled_reason = 'manual'
|
||||
WHERE is_active = 0 AND disabled_reason IS NULL
|
||||
""")
|
||||
# For existing expired tokens, set to 'expired'
|
||||
await db.execute("""
|
||||
UPDATE tokens
|
||||
SET disabled_reason = 'expired'
|
||||
WHERE is_expired = 1 AND disabled_reason IS NULL
|
||||
""")
|
||||
|
||||
await db.commit()
|
||||
|
||||
async def init_config_from_toml(self, config_dict: dict, is_first_startup: bool = True):
|
||||
@@ -698,27 +715,35 @@ class Database:
|
||||
""", (token_id,))
|
||||
await db.commit()
|
||||
|
||||
async def update_token_status(self, token_id: int, is_active: bool):
|
||||
"""Update token status"""
|
||||
async def update_token_status(self, token_id: int, is_active: bool, disabled_reason: Optional[str] = None):
|
||||
"""Update token status and disabled reason"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("""
|
||||
UPDATE tokens SET is_active = ? WHERE id = ?
|
||||
""", (is_active, token_id))
|
||||
UPDATE tokens SET is_active = ?, disabled_reason = ? WHERE id = ?
|
||||
""", (is_active, disabled_reason, token_id))
|
||||
await db.commit()
|
||||
|
||||
async def mark_token_expired(self, token_id: int):
|
||||
"""Mark token as expired and disable it"""
|
||||
"""Mark token as expired and disable it with reason"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("""
|
||||
UPDATE tokens SET is_expired = 1, is_active = 0 WHERE id = ?
|
||||
""", (token_id,))
|
||||
UPDATE tokens SET is_expired = 1, is_active = 0, disabled_reason = ? WHERE id = ?
|
||||
""", ("expired", token_id))
|
||||
await db.commit()
|
||||
|
||||
async def mark_token_invalid(self, token_id: int):
|
||||
"""Mark token as invalid (401 error) and disable it"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("""
|
||||
UPDATE tokens SET is_expired = 1, is_active = 0, disabled_reason = ? WHERE id = ?
|
||||
""", ("token_invalid", token_id))
|
||||
await db.commit()
|
||||
|
||||
async def clear_token_expired(self, token_id: int):
|
||||
"""Clear token expired flag"""
|
||||
"""Clear token expired flag and disabled reason"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("""
|
||||
UPDATE tokens SET is_expired = 0 WHERE id = ?
|
||||
UPDATE tokens SET is_expired = 0, disabled_reason = NULL WHERE id = ?
|
||||
""", (token_id,))
|
||||
await db.commit()
|
||||
|
||||
|
||||
@@ -40,6 +40,8 @@ class Token(BaseModel):
|
||||
video_concurrency: int = -1 # 视频并发数限制,-1表示不限制
|
||||
# 过期标记
|
||||
is_expired: bool = False # Token是否已过期(401 token_invalidated)
|
||||
# 禁用原因: manual=手动禁用, error_limit=错误次数超限, token_invalid=Token失效, expired=过期失效
|
||||
disabled_reason: Optional[str] = None
|
||||
|
||||
class TokenStats(BaseModel):
|
||||
"""Token statistics"""
|
||||
|
||||
@@ -946,19 +946,21 @@ class TokenManager:
|
||||
|
||||
async def update_token_status(self, token_id: int, is_active: bool):
|
||||
"""Update token active status"""
|
||||
await self.db.update_token_status(token_id, is_active)
|
||||
# When manually changing status, set appropriate disabled_reason
|
||||
disabled_reason = None if is_active else "manual"
|
||||
await self.db.update_token_status(token_id, is_active, disabled_reason)
|
||||
|
||||
async def enable_token(self, token_id: int):
|
||||
"""Enable a token and reset error count"""
|
||||
await self.db.update_token_status(token_id, True)
|
||||
await self.db.update_token_status(token_id, True, None) # Clear disabled_reason
|
||||
# Reset error count when enabling (in token_stats table)
|
||||
await self.db.reset_error_count(token_id)
|
||||
# Clear expired flag when enabling
|
||||
await self.db.clear_token_expired(token_id)
|
||||
|
||||
async def disable_token(self, token_id: int):
|
||||
"""Disable a token"""
|
||||
await self.db.update_token_status(token_id, False)
|
||||
"""Disable a token (manual disable)"""
|
||||
await self.db.update_token_status(token_id, False, "manual")
|
||||
|
||||
async def test_token(self, token_id: int) -> dict:
|
||||
"""Test if a token is valid by calling Sora API and refresh account info (subscription + Sora2)"""
|
||||
@@ -1048,6 +1050,14 @@ class TokenManager:
|
||||
"valid": False,
|
||||
"message": "Token已过期(token_invalidated)"
|
||||
}
|
||||
# Check if error is "Failed to get user info:401"
|
||||
if "Failed to get user info:401" in error_msg or "Failed to get user info: 401" in error_msg:
|
||||
# Mark token as invalid and disable it
|
||||
await self.db.mark_token_invalid(token_id)
|
||||
return {
|
||||
"valid": False,
|
||||
"message": "Token无效: Token is invalid: Failed to get user info:401"
|
||||
}
|
||||
return {
|
||||
"valid": False,
|
||||
"message": f"Token is invalid: {error_msg}"
|
||||
@@ -1077,7 +1087,8 @@ class TokenManager:
|
||||
admin_config = await self.db.get_admin_config()
|
||||
|
||||
if stats and stats.consecutive_error_count >= admin_config.error_ban_threshold:
|
||||
await self.db.update_token_status(token_id, False)
|
||||
# Disable token with error_limit reason
|
||||
await self.db.update_token_status(token_id, False, "error_limit")
|
||||
|
||||
async def record_success(self, token_id: int, is_video: bool = False):
|
||||
"""Record successful request (reset error count)"""
|
||||
|
||||
Reference in New Issue
Block a user