Spaces:
Paused
Paused
| """ | |
| UI State Management | |
| """ | |
| import asyncio | |
| import json | |
| import logging | |
| from playwright.async_api import Page as AsyncPage | |
| from logging_utils import set_request_id | |
| logger = logging.getLogger("AIStudioProxyServer") | |
| async def _verify_ui_state_settings(page: AsyncPage, req_id: str = "unknown") -> dict: | |
| """ | |
| Verify if the UI state settings are correct. | |
| Args: | |
| page: Playwright page object. | |
| req_id: Request ID for logging. | |
| Returns: | |
| dict: A dictionary containing the validation result. | |
| """ | |
| # Don't set lifecycle phase names as request IDs - they appear as ghost prefixes | |
| # Only set actual request IDs (7-char alphanumerics) | |
| if req_id not in ("initial", "set_mod", "set_model", "reload", "unknown", ""): | |
| set_request_id(req_id) | |
| try: | |
| logger.debug("[State] Verifying UI state...") | |
| # Get current localStorage settings | |
| prefs_str = await page.evaluate( | |
| "() => localStorage.getItem('aiStudioUserPreference')" | |
| ) | |
| if not prefs_str: | |
| logger.warning("localStorage.aiStudioUserPreference does not exist") | |
| return { | |
| "exists": False, | |
| "isAdvancedOpen": None, | |
| "areToolsOpen": None, | |
| "needsUpdate": True, | |
| "error": "localStorage not found", | |
| } | |
| try: | |
| prefs = json.loads(prefs_str) | |
| is_advanced_open = prefs.get("isAdvancedOpen") | |
| are_tools_open = prefs.get("areToolsOpen") | |
| # Check if update is needed | |
| needs_update = (is_advanced_open is not True) or ( | |
| are_tools_open is not True | |
| ) | |
| result = { | |
| "exists": True, | |
| "isAdvancedOpen": is_advanced_open, | |
| "areToolsOpen": are_tools_open, | |
| "needsUpdate": needs_update, | |
| "prefs": prefs, | |
| } | |
| if needs_update: | |
| logger.debug( | |
| f"[State] State mismatch: adv={is_advanced_open}, tools={are_tools_open} (update needed)" | |
| ) | |
| # No log needed when state is correct | |
| return result | |
| except json.JSONDecodeError as e: | |
| logger.error(f"Failed to parse localStorage JSON: {e}") | |
| return { | |
| "exists": False, | |
| "isAdvancedOpen": None, | |
| "areToolsOpen": None, | |
| "needsUpdate": True, | |
| "error": f"JSON parse failed: {e}", | |
| } | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as e: | |
| logger.error(f"Error verifying UI state settings: {e}") | |
| return { | |
| "exists": False, | |
| "isAdvancedOpen": None, | |
| "areToolsOpen": None, | |
| "needsUpdate": True, | |
| "error": f"Verification failed: {e}", | |
| } | |
| async def _force_ui_state_settings(page: AsyncPage, req_id: str = "unknown") -> bool: | |
| """ | |
| Forcefully set the UI state. | |
| Args: | |
| page: Playwright page object. | |
| req_id: Request ID for logging. | |
| Returns: | |
| bool: Whether the setting was successful. | |
| """ | |
| try: | |
| logger.debug("[State] Forcefully setting UI state...") | |
| # First verify current state | |
| current_state = await _verify_ui_state_settings(page, req_id) | |
| if not current_state["needsUpdate"]: | |
| logger.debug("[State] State is already correct, no update needed") | |
| return True | |
| # Get existing preferences or create new ones | |
| prefs = current_state.get("prefs", {}) | |
| # Force key configurations | |
| prefs["isAdvancedOpen"] = True | |
| prefs["areToolsOpen"] = True | |
| # Save to localStorage | |
| prefs_str = json.dumps(prefs) | |
| await page.evaluate( | |
| "(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)", | |
| prefs_str, | |
| ) | |
| logger.debug("[State] Set: isAdvancedOpen=true, areToolsOpen=true") | |
| # Verify if setting was successful | |
| verify_state = await _verify_ui_state_settings(page, req_id) | |
| if not verify_state["needsUpdate"]: | |
| logger.debug("[State] Setting verification successful") | |
| return True | |
| else: | |
| logger.warning("UI state setting verification failed, may need to retry") | |
| return False | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as e: | |
| logger.error(f"Error forcefully setting UI state: {e}") | |
| return False | |
| async def _force_ui_state_with_retry( | |
| page: AsyncPage, | |
| req_id: str = "unknown", | |
| max_retries: int = 3, | |
| retry_delay: float = 1.0, | |
| ) -> bool: | |
| """ | |
| Forcefully set UI state with retry mechanism. | |
| Args: | |
| page: Playwright page object. | |
| req_id: Request ID for logging. | |
| max_retries: Maximum number of retries. | |
| retry_delay: Retry delay (seconds). | |
| Returns: | |
| bool: Whether the setting was ultimately successful. | |
| """ | |
| for attempt in range(1, max_retries + 1): | |
| success = await _force_ui_state_settings(page, req_id) | |
| if success: | |
| return True | |
| if attempt < max_retries: | |
| logger.debug(f"[State] Retrying {attempt}/{max_retries}...") | |
| await asyncio.sleep(retry_delay) | |
| else: | |
| logger.warning(f"[State] Still failed after {max_retries} attempts") | |
| return False | |
| async def _verify_and_apply_ui_state(page: AsyncPage, req_id: str = "unknown") -> bool: | |
| """ | |
| Full process of verifying and applying UI state settings. | |
| Args: | |
| page: Playwright page object. | |
| req_id: Request ID for logging. | |
| Returns: | |
| bool: Whether the operation was successful. | |
| """ | |
| try: | |
| logger.debug("[State] Starting to verify and apply UI state...") | |
| # First verify current state | |
| state = await _verify_ui_state_settings(page, req_id) | |
| if state["needsUpdate"]: | |
| logger.debug("[State] Update needed, applying forced settings...") | |
| return await _force_ui_state_with_retry(page, req_id) | |
| else: | |
| return True | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as e: | |
| logger.error(f"Error during verifying and applying UI state: {e}") | |
| return False | |