Spaces:
Paused
Paused
File size: 6,412 Bytes
6480add af689c4 6480add af689c4 b7dd118 6480add |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
"""
Cookie pool manager for Z.AI tokens with round-robin rotation
"""
import asyncio
import logging
from typing import List, Optional
from asyncio import Lock
import httpx
from config import settings
logger = logging.getLogger(__name__)
class CookieManager:
def __init__(self, cookies: List[str]):
self.cookies = cookies or []
self.current_index = 0
self.lock = Lock()
self.failed_cookies = set()
if self.cookies:
logger.info(f"Initialized CookieManager with {len(cookies)} cookies")
else:
logger.warning("CookieManager initialized with no cookies")
async def get_next_cookie(self) -> Optional[str]:
"""Get the next available cookie using round-robin"""
if not self.cookies:
return None
async with self.lock:
attempts = 0
while attempts < len(self.cookies):
cookie = self.cookies[self.current_index]
self.current_index = (self.current_index + 1) % len(self.cookies)
# Skip failed cookies
if cookie not in self.failed_cookies:
return cookie
attempts += 1
# All cookies failed, reset failed set and try again
if self.failed_cookies:
logger.warning(f"All {len(self.cookies)} cookies failed, resetting failed set and retrying")
self.failed_cookies.clear()
return self.cookies[0]
return None
async def mark_cookie_failed(self, cookie: str):
"""Mark a cookie as failed"""
async with self.lock:
self.failed_cookies.add(cookie)
logger.warning(f"Marked cookie as failed: {cookie[:20]}...")
async def mark_cookie_success(self, cookie: str):
"""Mark a cookie as working (remove from failed set)"""
async with self.lock:
if cookie in self.failed_cookies:
self.failed_cookies.discard(cookie)
logger.info(f"Cookie recovered: {cookie[:20]}...")
async def health_check(self, cookie: str) -> bool:
"""Check if a cookie is still valid"""
try:
async with httpx.AsyncClient() as client:
# Use the same payload format as actual requests
import uuid
test_payload = {
"stream": True,
"model": "GLM-4-6-API-V1",
"messages": [{"role": "user", "content": "hi"}],
"background_tasks": {
"title_generation": False,
"tags_generation": False
},
"chat_id": str(uuid.uuid4()),
"features": {
"image_generation": False,
"code_interpreter": False,
"web_search": False,
"auto_web_search": False
},
"id": str(uuid.uuid4()),
"mcp_servers": [],
"model_item": {
"id": "GLM-4-6-API-V1",
"name": "GLM-4.6",
"owned_by": "openai"
},
"params": {},
"tool_servers": [],
"variables": {
"{{USER_NAME}}": "User",
"{{USER_LOCATION}}": "Unknown",
"{{CURRENT_DATETIME}}": "2025-08-04 16:46:56"
}
}
response = await client.post(
"https://chat.z.ai/api/chat/completions",
headers={
"Authorization": f"Bearer {cookie}",
"Content-Type": "application/json",
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36",
"Accept": "application/json, text/event-stream",
"Accept-Language": "zh-CN",
"sec-ch-ua": '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"',
"sec-ch-ua-mobile": "?0",
"sec-ch-ua-platform": '"macOS"',
"x-fe-version": "prod-fe-1.0.53",
"Origin": "https://chat.z.ai",
"Referer": "https://chat.z.ai/c/069723d5-060b-404f-992c-4705f1554c4c"
},
json=test_payload,
timeout=10.0
)
# Consider 200 as success
is_healthy = response.status_code == 200
if not is_healthy:
logger.debug(f"Health check failed for cookie {cookie[:20]}...: HTTP {response.status_code}")
else:
logger.debug(f"Health check passed for cookie {cookie[:20]}...")
return is_healthy
except Exception as e:
logger.debug(f"Health check failed for cookie {cookie[:20]}...: {e}")
return False
async def periodic_health_check(self):
"""Periodically check all cookies health"""
while True:
try:
# Only check if we have cookies and some are marked as failed
if self.cookies and self.failed_cookies:
logger.info(f"Running health check for {len(self.failed_cookies)} failed cookies")
for cookie in list(self.failed_cookies): # Create a copy to avoid modification during iteration
if await self.health_check(cookie):
await self.mark_cookie_success(cookie)
logger.info(f"Cookie recovered: {cookie[:20]}...")
else:
logger.debug(f"Cookie still failed: {cookie[:20]}...")
# Wait 10 minutes before next check (reduced frequency)
await asyncio.sleep(600)
except Exception as e:
logger.error(f"Error in periodic health check: {e}")
await asyncio.sleep(300) # Wait 5 minutes on error
# Global cookie manager instance
cookie_manager = CookieManager(settings.COOKIES if settings else [])
|