Spaces:
Paused
Paused
File size: 2,365 Bytes
a5784e9 | 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 | import asyncio
import logging
import os
from dotenv import load_dotenv
load_dotenv()
# --- Centralized state module ---
from api_utils.server_state import state
def clear_debug_logs() -> None:
state.clear_debug_logs()
# --- Imports ---
from browser_utils.auth_rotation import perform_auth_rotation
from config import (
GlobalState,
)
async def quota_watchdog():
"""Background watchdog to monitor quota exceeded events."""
# Use state's logger if available
logger = getattr(state, "logger", logging.getLogger("AIStudioProxyServer"))
logger.info("👀 Quota Watchdog Started")
while True:
try:
await GlobalState.QUOTA_EXCEEDED_EVENT.wait()
logger.critical(
"🚨 Watchdog detected Quota Exceeded! Initiating Rotation..."
)
if not GlobalState.AUTH_ROTATION_LOCK.is_set():
logger.info("Watchdog: Rotation already in progress. Waiting...")
await asyncio.sleep(1)
continue
GlobalState.start_recovery()
try:
current_model_id = state.current_ai_studio_model_id
success = await perform_auth_rotation(
target_model_id=current_model_id or ""
)
if success:
logger.info("Watchdog: Rotation successful.")
else:
logger.error("Watchdog: Rotation failed.")
finally:
GlobalState.finish_recovery()
if GlobalState.IS_QUOTA_EXCEEDED:
logger.warning("Watchdog: Quota flag still set. Forcing reset.")
GlobalState.reset_quota_status()
except asyncio.CancelledError:
logger.info("Watchdog: Task cancelled.")
break
except Exception as e:
logger.error(f"Watchdog Error: {e}", exc_info=True)
await asyncio.sleep(5)
# Register quota_watchdog in state for easier access and to avoid circular import issues
state.quota_watchdog = quota_watchdog
from api_utils import (
create_app,
)
# --- FastAPI App ---
app = create_app()
if __name__ == "__main__":
import uvicorn
port = int(os.environ.get("PORT", 2048))
uvicorn.run(
"server:app", host="0.0.0.0", port=port, log_level="info", access_log=False
)
|