Spaces:
Paused
Paused
File size: 4,368 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 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 | """
Centralized server state module.
This module contains all shared state variables that were previously in server.py.
It has NO imports from project modules (only stdlib), making it safe to import
at module level anywhere in the codebase without circular dependency issues.
Usage:
from api_utils.server_state import state
# Access state attributes
page = state.page_instance
state.current_ai_studio_model_id = "new-model"
"""
import asyncio
import logging
import multiprocessing
from asyncio import Event, Lock, Queue, Task
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set
if TYPE_CHECKING:
from playwright.async_api import (
Browser as AsyncBrowser,
)
from playwright.async_api import (
Page as AsyncPage,
)
from playwright.async_api import (
Playwright as AsyncPlaywright,
)
from api_utils.context_types import QueueItem
from models.logging import WebSocketConnectionManager
class ServerState:
"""
Centralized container for all server state.
This class holds all mutable state that needs to be shared across modules.
Using a class allows for better organization and easier testing (state can be reset).
"""
def __init__(self) -> None:
"""Initialize all state variables with default values."""
self.reset()
def reset(self) -> None:
"""Reset all state to initial values. Useful for testing."""
# --- Stream Queue ---
self.STREAM_QUEUE: Optional[multiprocessing.Queue] = None
self.STREAM_PROCESS: Optional[multiprocessing.Process] = None
# --- Playwright/Browser State ---
self.playwright_manager: Optional["AsyncPlaywright"] = None
self.browser_instance: Optional["AsyncBrowser"] = None
self.page_instance: Optional["AsyncPage"] = None
self.is_playwright_ready: bool = False
self.is_browser_connected: bool = False
self.is_page_ready: bool = False
self.is_initializing: bool = False
# --- Proxy Configuration ---
self.PLAYWRIGHT_PROXY_SETTINGS: Optional[Dict[str, str]] = None
# --- Model State ---
self.global_model_list_raw_json: Optional[str] = None
self.parsed_model_list: List[Dict[str, Any]] = []
self.model_list_fetch_event: Event = asyncio.Event()
self.current_ai_studio_model_id: Optional[str] = None
self.current_auth_profile_path: Optional[str] = None
self.model_switching_lock: Lock = Lock()
self.excluded_model_ids: Set[str] = set()
# --- Request Processing State ---
self.request_queue: "Optional[Queue[QueueItem]]" = None
self.processing_lock: Optional[Lock] = None
self.worker_task: "Optional[Task[None]]" = None
# --- Parameter Cache ---
self.page_params_cache: Dict[str, Any] = {}
self.params_cache_lock: Lock = Lock()
# --- Debug Logging State ---
self.console_logs: List[Dict[str, Any]] = []
self.network_log: Dict[str, List[Dict[str, Any]]] = {
"requests": [],
"responses": [],
}
# --- Logging ---
self.logger: logging.Logger = logging.getLogger("AIStudioProxyServer")
self.log_ws_manager: Optional["WebSocketConnectionManager"] = None
# --- Control Flags ---
self.should_exit: bool = False
self.quota_watchdog: Optional[Callable] = None
def clear_debug_logs(self) -> None:
"""Clear console and network logs (called after each request)."""
self.console_logs = []
self.network_log = {"requests": [], "responses": []}
# Global singleton instance
state = ServerState()
# Convenience exports for backward compatibility
# These allow direct attribute access like: from server_state import page_instance
# But the recommended way is: from api_utils.server_state import state; state.page_instance
def __getattr__(name: str) -> Any:
"""
Module-level attribute access for backward compatibility.
Allows:
from server_state import page_instance
Instead of:
from api_utils.server_state import state
page_instance = state.page_instance
"""
if hasattr(state, name):
return getattr(state, name)
raise AttributeError(f"module 'server_state' has no attribute '{name}'")
|