Spaces:
Paused
Paused
File size: 10,914 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 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 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 | """
Model State Initialization and Synchronization
"""
import asyncio
import json
import logging
from playwright.async_api import Page as AsyncPage
from playwright.async_api import expect as expect_async
from config import INPUT_SELECTOR, MODEL_NAME_SELECTOR
from .ui_state import _verify_and_apply_ui_state, _verify_ui_state_settings
logger = logging.getLogger("AIStudioProxyServer")
async def _handle_initial_model_state_and_storage(page: AsyncPage):
"""Handle initial model state and storage"""
from api_utils.server_state import state
getattr(state, "current_ai_studio_model_id", None)
getattr(state, "parsed_model_list", [])
getattr(state, "model_list_fetch_event", None)
logger.debug("[Init] Processing initial model state and localStorage...")
needs_reload_and_storage_update = False
reason_for_reload = ""
try:
initial_prefs_str = await page.evaluate(
"() => localStorage.getItem('aiStudioUserPreference')"
)
if not initial_prefs_str:
needs_reload_and_storage_update = True
reason_for_reload = "localStorage not found"
else:
try:
pref_obj = json.loads(initial_prefs_str)
prompt_model_path = pref_obj.get("promptModel")
pref_obj.get("isAdvancedOpen")
is_prompt_model_valid = (
isinstance(prompt_model_path, str) and prompt_model_path.strip()
)
if not is_prompt_model_valid:
needs_reload_and_storage_update = True
reason_for_reload = "promptModel invalid"
else:
# Use new UI state verification
ui_state = await _verify_ui_state_settings(page, "initial")
if ui_state["needsUpdate"]:
needs_reload_and_storage_update = True
reason_for_reload = "UI state mismatch"
else:
state.current_ai_studio_model_id = prompt_model_path.split("/")[
-1
]
logger.debug(
f"localStorage valid and UI state correct. Initial model ID set from localStorage: {state.current_ai_studio_model_id}"
)
except json.JSONDecodeError:
needs_reload_and_storage_update = True
reason_for_reload = (
"Failed to parse localStorage.aiStudioUserPreference JSON."
)
logger.error(
f"Determined refresh and storage update needed: {reason_for_reload}"
)
if needs_reload_and_storage_update:
logger.debug(f"[State] Refresh needed: {reason_for_reload}")
await _set_model_from_page_display(page, set_storage=True)
current_page_url = page.url
logger.info("[UI Operation] Reloading page to apply settings...")
max_retries = 3
for attempt in range(max_retries):
try:
logger.debug(
f"Attempting page reload (attempt {attempt + 1}/{max_retries}): {current_page_url}"
)
await page.goto(
current_page_url, wait_until="domcontentloaded", timeout=40000
)
await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
timeout=30000
)
logger.debug(f"Page successfully reloaded to: {page.url}")
# Verify UI state after page reload
logger.debug("[State] Verifying UI state...")
reload_ui_state_success = await _verify_and_apply_ui_state(
page, "reload"
)
if reload_ui_state_success:
logger.info("[UI Check] Verification passed after page reload")
else:
logger.warning("UI state verification failed after reload")
break # Exit loop on success
except asyncio.CancelledError:
raise
except Exception as reload_err:
logger.warning(
f"Page reload attempt {attempt + 1}/{max_retries} failed: {reload_err}"
)
if attempt < max_retries - 1:
logger.debug("[Init] Retrying in 5 seconds...")
await asyncio.sleep(5)
else:
logger.error(
f"Page reload ultimately failed after {max_retries} attempts: {reload_err}. Subsequent model state may be inaccurate.",
exc_info=True,
)
from browser_utils.operations import save_error_snapshot
await save_error_snapshot(
f"initial_storage_reload_fail_attempt_{attempt + 1}"
)
logger.debug("[State] Syncing model ID after reload")
await _set_model_from_page_display(page, set_storage=False)
logger.debug(
f"[State] Complete, current model: {state.current_ai_studio_model_id}"
)
else:
logger.debug("[State] localStorage state OK, no refresh needed")
except asyncio.CancelledError:
raise
except Exception as e:
logger.error(
f"(New) Critical error processing initial model state and localStorage: {e}",
exc_info=True,
)
try:
logger.warning(
"Due to error, attempting fallback to set global model ID from page display only (not writing to localStorage)..."
)
await _set_model_from_page_display(page, set_storage=False)
except asyncio.CancelledError:
raise
except Exception as fallback_err:
logger.error(f"Fallback model ID setting also failed: {fallback_err}")
async def _set_model_from_page_display(page: AsyncPage, set_storage: bool = False):
"""Set model from page display"""
from api_utils.server_state import state
getattr(state, "current_ai_studio_model_id", None)
getattr(state, "parsed_model_list", [])
model_list_fetch_event = getattr(state, "model_list_fetch_event", None)
try:
logger.debug("[Model] Reading current model from page display...")
model_name_locator = page.locator(MODEL_NAME_SELECTOR)
displayed_model_name_from_page_raw = await model_name_locator.first.inner_text(
timeout=7000
)
displayed_model_name = displayed_model_name_from_page_raw.strip()
logger.debug(f"[Model] Page display: '{displayed_model_name}'")
found_model_id_from_display = None
if model_list_fetch_event and not model_list_fetch_event.is_set():
logger.debug("[Model] Waiting for model list data (up to 5 seconds)...")
try:
await asyncio.wait_for(model_list_fetch_event.wait(), timeout=5.0)
except asyncio.TimeoutError:
logger.warning(
"Timeout waiting for model list, may not be able to accurately convert display name to ID."
)
found_model_id_from_display = displayed_model_name
new_model_value = found_model_id_from_display
if state.current_ai_studio_model_id != new_model_value:
state.current_ai_studio_model_id = new_model_value
logger.debug(f"[Model] Global ID updated: {new_model_value}")
# No log needed if unchanged
if set_storage:
logger.debug("[State] Preparing to update localStorage")
existing_prefs_for_update_str = await page.evaluate(
"() => localStorage.getItem('aiStudioUserPreference')"
)
prefs_to_set = {}
if existing_prefs_for_update_str:
try:
prefs_to_set = json.loads(existing_prefs_for_update_str)
except json.JSONDecodeError:
logger.warning(
"Failed to parse existing localStorage.aiStudioUserPreference, will create new preferences."
)
# Use new force settings feature
logger.debug("[State] Applying forced UI state settings...")
ui_state_success = await _verify_and_apply_ui_state(page, "set_model")
if not ui_state_success:
logger.warning("UI state setting failed, using legacy method")
prefs_to_set["isAdvancedOpen"] = True
prefs_to_set["areToolsOpen"] = True
else:
# Ensure prefs_to_set also contains correct settings
prefs_to_set["isAdvancedOpen"] = True
prefs_to_set["areToolsOpen"] = True
logger.debug("[State] Set: isAdvancedOpen=true, areToolsOpen=true")
if found_model_id_from_display:
new_prompt_model_path = f"models/{found_model_id_from_display}"
prefs_to_set["promptModel"] = new_prompt_model_path
elif "promptModel" not in prefs_to_set:
logger.warning(
f"Could not find model ID from page display '{displayed_model_name}', and no existing promptModel in localStorage. promptModel will not be actively set to avoid potential issues."
)
default_keys_if_missing = {
"bidiModel": "models/gemini-1.0-pro-001",
"isSafetySettingsOpen": False,
"hasShownSearchGroundingTos": False,
"autosaveEnabled": True,
"theme": "system",
"bidiOutputFormat": 3,
"isSystemInstructionsOpen": False,
"warmWelcomeDisplayed": True,
"getCodeLanguage": "Node.js",
"getCodeHistoryToggle": False,
"fileCopyrightAcknowledged": True,
}
for key, val_default in default_keys_if_missing.items():
if key not in prefs_to_set:
prefs_to_set[key] = val_default
await page.evaluate(
"(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)",
json.dumps(prefs_to_set),
)
logger.debug(
f"[State] localStorage updated (model: {prefs_to_set.get('promptModel', 'N/A')})"
)
except asyncio.CancelledError:
raise
except Exception as e_set_disp:
logger.error(
f"Error setting model from page display: {e_set_disp}", exc_info=True
)
|