Function quota_watchdog

server.py:25–64

Background watchdog to monitor quota exceeded events.

()

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