browser_utils/model_management.py

# --- browser_utils/model_management.py ---
# 浏览器模型管理相关功能模块

import asyncio
import json
import os
import logging
import time
from typing import Optional, Set

from playwright.async_api import Page as AsyncPage, expect as expect_async, Error as PlaywrightAsyncError

# 导入配置和模型
from config import *
from models import ClientDisconnectedError

logger = logging.getLogger("AIStudioProxyServer")

# ==================== 强制UI状态设置功能 ====================

async def _verify_ui_state_settings(page: AsyncPage, req_id: str = "unknown") -> dict:
    """
    验证UI状态设置是否正确

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志

    Returns:
        dict: 包含验证结果的字典
    """
    try:
        logger.info(f"[{req_id}] 验证UI状态设置...")

        # 获取当前localStorage设置
        prefs_str = await page.evaluate("() => localStorage.getItem('aiStudioUserPreference')")

        if not prefs_str:
            logger.warning(f"[{req_id}] localStorage.aiStudioUserPreference 不存在")
            return {
                'exists': False,
                'isAdvancedOpen': None,
                'areToolsOpen': None,
                'needsUpdate': True,
                'error': 'localStorage不存在'
            }

        try:
            prefs = json.loads(prefs_str)
            is_advanced_open = prefs.get('isAdvancedOpen')
            are_tools_open = prefs.get('areToolsOpen')

            # 检查是否需要更新
            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
            }

            logger.info(f"[{req_id}] UI状态验证结果: isAdvancedOpen={is_advanced_open}, areToolsOpen={are_tools_open} (期望: True), needsUpdate={needs_update}")
            return result

        except json.JSONDecodeError as e:
            logger.error(f"[{req_id}] 解析localStorage JSON失败: {e}")
            return {
                'exists': False,
                'isAdvancedOpen': None,
                'areToolsOpen': None,
                'needsUpdate': True,
                'error': f'JSON解析失败: {e}'
            }

    except Exception as e:
        logger.error(f"[{req_id}] 验证UI状态设置时发生错误: {e}")
        return {
            'exists': False,
            'isAdvancedOpen': None,
            'areToolsOpen': None,
            'needsUpdate': True,
            'error': f'验证失败: {e}'
        }

async def _force_ui_state_settings(page: AsyncPage, req_id: str = "unknown") -> bool:
    """
    强制设置UI状态

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志

    Returns:
        bool: 设置是否成功
    """
    try:
        logger.info(f"[{req_id}] 开始强制设置UI状态...")

        # 首先验证当前状态
        current_state = await _verify_ui_state_settings(page, req_id)

        if not current_state['needsUpdate']:
            logger.info(f"[{req_id}] UI状态已正确设置，无需更新")
            return True

        # 获取现有preferences或创建新的
        prefs = current_state.get('prefs', {})

        # 强制设置关键配置
        prefs['isAdvancedOpen'] = True
        prefs['areToolsOpen'] = True

        # 保存到localStorage
        prefs_str = json.dumps(prefs)
        await page.evaluate("(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)", prefs_str)

        logger.info(f"[{req_id}] 已强制设置: isAdvancedOpen=true, areToolsOpen=true")

        # 验证设置是否成功
        verify_state = await _verify_ui_state_settings(page, req_id)
        if not verify_state['needsUpdate']:
            logger.info(f"[{req_id}] ✅ UI状态设置验证成功")
            return True
        else:
            logger.warning(f"[{req_id}] ⚠️ UI状态设置验证失败，可能需要重试")
            return False

    except Exception as e:
        logger.error(f"[{req_id}] 强制设置UI状态时发生错误: {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:
    """
    带重试机制的UI状态强制设置

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志
        max_retries: 最大重试次数
        retry_delay: 重试延迟（秒）

    Returns:
        bool: 设置是否最终成功
    """
    for attempt in range(1, max_retries + 1):
        logger.info(f"[{req_id}] 尝试强制设置UI状态 (第 {attempt}/{max_retries} 次)")

        success = await _force_ui_state_settings(page, req_id)
        if success:
            logger.info(f"[{req_id}] ✅ UI状态设置在第 {attempt} 次尝试中成功")
            return True

        if attempt < max_retries:
            logger.warning(f"[{req_id}] ⚠️ 第 {attempt} 次尝试失败，{retry_delay}秒后重试...")
            await asyncio.sleep(retry_delay)
        else:
            logger.error(f"[{req_id}] ❌ UI状态设置在 {max_retries} 次尝试后仍然失败")

    return False

async def _verify_and_apply_ui_state(page: AsyncPage, req_id: str = "unknown") -> bool:
    """
    验证并应用UI状态设置的完整流程

    Args:
        page: Playwright页面对象
        req_id: 请求ID用于日志

    Returns:
        bool: 操作是否成功
    """
    try:
        logger.info(f"[{req_id}] 开始验证并应用UI状态设置...")

        # 首先验证当前状态
        state = await _verify_ui_state_settings(page, req_id)

        logger.info(f"[{req_id}] 当前UI状态: exists={state['exists']}, isAdvancedOpen={state['isAdvancedOpen']}, areToolsOpen={state['areToolsOpen']}, needsUpdate={state['needsUpdate']}")

        if state['needsUpdate']:
            logger.info(f"[{req_id}] 检测到UI状态需要更新，正在应用强制设置...")
            return await _force_ui_state_with_retry(page, req_id)
        else:
            logger.info(f"[{req_id}] UI状态已正确设置，无需更新")
            return True

    except Exception as e:
        logger.error(f"[{req_id}] 验证并应用UI状态设置时发生错误: {e}")
        return False

async def switch_ai_studio_model(page: AsyncPage, model_id: str, req_id: str) -> bool:
    """切换AI Studio模型"""
    logger.info(f"[{req_id}] 开始切换模型到: {model_id}")
    original_prefs_str: Optional[str] = None
    original_prompt_model: Optional[str] = None
    new_chat_url = f"https://{AI_STUDIO_URL_PATTERN}prompts/new_chat"
    
    try:
        original_prefs_str = await page.evaluate("() => localStorage.getItem('aiStudioUserPreference')")
        if original_prefs_str:
            try:
                original_prefs_obj = json.loads(original_prefs_str)
                original_prompt_model = original_prefs_obj.get("promptModel")
                logger.info(f"[{req_id}] 切换前 localStorage.promptModel 为: {original_prompt_model or '未设置'}")
            except json.JSONDecodeError:
                logger.warning(f"[{req_id}] 无法解析原始的 aiStudioUserPreference JSON 字符串。")
                original_prefs_str = None
        
        current_prefs_for_modification = json.loads(original_prefs_str) if original_prefs_str else {}
        full_model_path = f"models/{model_id}"
        
        if current_prefs_for_modification.get("promptModel") == full_model_path:
            logger.info(f"[{req_id}] 模型已经设置为 {model_id} (localStorage 中已是目标值)，无需切换")
            if page.url != new_chat_url:
                 logger.info(f"[{req_id}] 当前 URL 不是 new_chat ({page.url})，导航到 {new_chat_url}")
                 await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=30000)
                 await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(timeout=30000)
            return True
        
        logger.info(f"[{req_id}] 从 {current_prefs_for_modification.get('promptModel', '未知')} 更新 localStorage.promptModel 为 {full_model_path}")
        current_prefs_for_modification["promptModel"] = full_model_path
        await page.evaluate("(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)", json.dumps(current_prefs_for_modification))
        
        # 使用新的强制设置功能
        logger.info(f"[{req_id}] 应用强制UI状态设置...")
        ui_state_success = await _verify_and_apply_ui_state(page, req_id)
        if not ui_state_success:
            logger.warning(f"[{req_id}] UI状态设置失败，但继续执行模型切换流程")

        # 为了保持兼容性，也更新当前的prefs对象
        current_prefs_for_modification["isAdvancedOpen"] = True
        current_prefs_for_modification["areToolsOpen"] = True
        await page.evaluate("(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)", json.dumps(current_prefs_for_modification))

        logger.info(f"[{req_id}] localStorage 已更新，导航到 '{new_chat_url}' 应用新模型...")
        await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=30000)

        input_field = page.locator(INPUT_SELECTOR)
        await expect_async(input_field).to_be_visible(timeout=30000)
        logger.info(f"[{req_id}] 页面已导航到新聊天并加载完成，输入框可见")

        # 页面加载后再次验证UI状态设置
        logger.info(f"[{req_id}] 页面加载完成，验证UI状态设置...")
        final_ui_state_success = await _verify_and_apply_ui_state(page, req_id)
        if final_ui_state_success:
            logger.info(f"[{req_id}] ✅ UI状态最终验证成功")
        else:
            logger.warning(f"[{req_id}] ⚠️ UI状态最终验证失败，但继续执行模型切换流程")
        
        final_prefs_str = await page.evaluate("() => localStorage.getItem('aiStudioUserPreference')")
        final_prompt_model_in_storage: Optional[str] = None
        if final_prefs_str:
            try:
                final_prefs_obj = json.loads(final_prefs_str)
                final_prompt_model_in_storage = final_prefs_obj.get("promptModel")
            except json.JSONDecodeError:
                logger.warning(f"[{req_id}] 无法解析刷新后的 aiStudioUserPreference JSON 字符串。")
        
        if final_prompt_model_in_storage == full_model_path:
            logger.info(f"[{req_id}] ✅ AI Studio localStorage 中模型已成功设置为: {full_model_path}")
            
            page_display_match = False
            expected_display_name_for_target_id = None
            actual_displayed_model_name_on_page = "无法读取"
            
            # 获取parsed_model_list
            import server
            parsed_model_list = getattr(server, 'parsed_model_list', [])
            
            if parsed_model_list:
                for m_obj in parsed_model_list:
                    if m_obj.get("id") == model_id:
                        expected_display_name_for_target_id = m_obj.get("display_name")
                        break

            try:
                model_name_locator = page.locator('[data-test-id="model-name"]')
                actual_displayed_model_id_on_page_raw = await model_name_locator.first.inner_text(timeout=5000)
                actual_displayed_model_id_on_page = actual_displayed_model_id_on_page_raw.strip()
                
                target_model_id = model_id

                if actual_displayed_model_id_on_page == target_model_id:
                    page_display_match = True
                    logger.info(f"[{req_id}] ✅ 页面显示模型ID ('{actual_displayed_model_id_on_page}') 与期望ID ('{target_model_id}') 一致。")
                else:
                    page_display_match = False
                    logger.error(f"[{req_id}] ❌ 页面显示模型ID ('{actual_displayed_model_id_on_page}') 与期望ID ('{target_model_id}') 不一致。")
            
            except Exception as e_disp:
                page_display_match = False # 读取失败则认为不匹配
                logger.warning(f"[{req_id}] 读取页面显示的当前模型ID时出错: {e_disp}。将无法验证页面显示。")

            if page_display_match:
                try:
                    logger.info(f"[{req_id}] 模型切换成功，重新启用 '临时聊天' 模式...")
                    incognito_button_locator = page.locator('button[aria-label="Temporary chat toggle"]')
                    
                    await incognito_button_locator.wait_for(state="visible", timeout=5000)
                    
                    button_classes = await incognito_button_locator.get_attribute("class")
                    
                    if button_classes and 'ms-button-active' in button_classes:
                        logger.info(f"[{req_id}] '临时聊天' 模式已处于激活状态。")
                    else:
                        logger.info(f"[{req_id}] '临时聊天' 模式未激活，正在点击以开启...")
                        await incognito_button_locator.click(timeout=3000)
                        await asyncio.sleep(0.5)
                        
                        updated_classes = await incognito_button_locator.get_attribute("class")
                        if updated_classes and 'ms-button-active' in updated_classes:
                             logger.info(f"[{req_id}] ✅ '临时聊天' 模式已成功重新启用。")
                        else:
                             logger.warning(f"[{req_id}] ⚠️ 点击后 '临时聊天' 模式状态验证失败，可能未成功重新开启。")
                
                except Exception as e:
                    logger.warning(f"[{req_id}] ⚠️ 模型切换后重新启用 '临时聊天' 模式失败: {e}")
                return True
            else:
                logger.error(f"[{req_id}] ❌ 模型切换失败，因为页面显示的模型与期望不符 (即使localStorage可能已更改)。")
        else:
            logger.error(f"[{req_id}] ❌ AI Studio 未接受模型更改 (localStorage)。期望='{full_model_path}', 实际='{final_prompt_model_in_storage or '未设置或无效'}'.")
        
        logger.info(f"[{req_id}] 模型切换失败。尝试恢复到页面当前实际显示的模型的状态...")
        current_displayed_name_for_revert_raw = "无法读取"
        current_displayed_name_for_revert_stripped = "无法读取"
        
        try:
            model_name_locator_revert = page.locator('[data-test-id="model-name"]')
            current_displayed_name_for_revert_raw = await model_name_locator_revert.first.inner_text(timeout=5000)
            current_displayed_name_for_revert_stripped = current_displayed_name_for_revert_raw.strip()
            logger.info(f"[{req_id}] 恢复：页面当前显示的模型名称 (原始: '{current_displayed_name_for_revert_raw}', 清理后: '{current_displayed_name_for_revert_stripped}')")
        except Exception as e_read_disp_revert:
            logger.warning(f"[{req_id}] 恢复：读取页面当前显示模型名称失败: {e_read_disp_revert}。将尝试回退到原始localStorage。")
            if original_prefs_str:
                logger.info(f"[{req_id}] 恢复：由于无法读取当前页面显示，尝试将 localStorage 恢复到原始状态: '{original_prompt_model or '未设置'}'")
                await page.evaluate("(origPrefs) => localStorage.setItem('aiStudioUserPreference', origPrefs)", original_prefs_str)
                logger.info(f"[{req_id}] 恢复：导航到 '{new_chat_url}' 以应用恢复的原始 localStorage 设置...")
                await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=20000)
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(timeout=20000)
                logger.info(f"[{req_id}] 恢复：页面已导航到新聊天并加载，已尝试应用原始 localStorage。")
            else:
                logger.warning(f"[{req_id}] 恢复：无有效的原始 localStorage 状态可恢复，也无法读取当前页面显示。")
            return False
        
        model_id_to_revert_to = None
        if current_displayed_name_for_revert_stripped != "无法读取":
            model_id_to_revert_to = current_displayed_name_for_revert_stripped
            logger.info(f"[{req_id}] 恢复：页面当前显示的ID是 '{model_id_to_revert_to}'，将直接用于恢复。")
        else:
            if current_displayed_name_for_revert_stripped == "无法读取":
                 logger.warning(f"[{req_id}] 恢复：因无法读取页面显示名称，故不能从 parsed_model_list 转换ID。")
            else:
                 logger.warning(f"[{req_id}] 恢复：parsed_model_list 为空，无法从显示名称 '{current_displayed_name_for_revert_stripped}' 转换模型ID。")
        
        if model_id_to_revert_to:
            base_prefs_for_final_revert = {}
            try:
                current_ls_content_str = await page.evaluate("() => localStorage.getItem('aiStudioUserPreference')")
                if current_ls_content_str:
                    base_prefs_for_final_revert = json.loads(current_ls_content_str)
                elif original_prefs_str:
                    base_prefs_for_final_revert = json.loads(original_prefs_str)
            except json.JSONDecodeError:
                logger.warning(f"[{req_id}] 恢复：解析现有 localStorage 以构建恢复偏好失败。")
            
            path_to_revert_to = f"models/{model_id_to_revert_to}"
            base_prefs_for_final_revert["promptModel"] = path_to_revert_to
            # 使用新的强制设置功能
            logger.info(f"[{req_id}] 恢复：应用强制UI状态设置...")
            ui_state_success = await _verify_and_apply_ui_state(page, req_id)
            if not ui_state_success:
                logger.warning(f"[{req_id}] 恢复：UI状态设置失败，但继续执行恢复流程")

            # 为了保持兼容性，也更新当前的prefs对象
            base_prefs_for_final_revert["isAdvancedOpen"] = True
            base_prefs_for_final_revert["areToolsOpen"] = True
            logger.info(f"[{req_id}] 恢复：准备将 localStorage.promptModel 设置回页面实际显示的模型的路径: '{path_to_revert_to}'，并强制设置配置选项")
            await page.evaluate("(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)", json.dumps(base_prefs_for_final_revert))
            logger.info(f"[{req_id}] 恢复：导航到 '{new_chat_url}' 以应用恢复到 '{model_id_to_revert_to}' 的 localStorage 设置...")
            await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=30000)
            await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(timeout=30000)

            # 恢复后再次验证UI状态
            logger.info(f"[{req_id}] 恢复：页面加载完成，验证UI状态设置...")
            final_ui_state_success = await _verify_and_apply_ui_state(page, req_id)
            if final_ui_state_success:
                logger.info(f"[{req_id}] ✅ 恢复：UI状态最终验证成功")
            else:
                logger.warning(f"[{req_id}] ⚠️ 恢复：UI状态最终验证失败")

            logger.info(f"[{req_id}] 恢复：页面已导航到新聊天并加载。localStorage 应已设置为反映模型 '{model_id_to_revert_to}'。")
        else:
            logger.error(f"[{req_id}] 恢复：无法将模型恢复到页面显示的状态，因为未能从显示名称 '{current_displayed_name_for_revert_stripped}' 确定有效模型ID。")
            if original_prefs_str:
                logger.warning(f"[{req_id}] 恢复：作为最终后备，尝试恢复到原始 localStorage: '{original_prompt_model or '未设置'}'")
                await page.evaluate("(origPrefs) => localStorage.setItem('aiStudioUserPreference', origPrefs)", original_prefs_str)
                logger.info(f"[{req_id}] 恢复：导航到 '{new_chat_url}' 以应用最终后备的原始 localStorage。")
                await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=20000)
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(timeout=20000)
                logger.info(f"[{req_id}] 恢复：页面已导航到新聊天并加载，已应用最终后备的原始 localStorage。")
            else:
                logger.warning(f"[{req_id}] 恢复：无有效的原始 localStorage 状态可作为最终后备。")
        
        return False
        
    except Exception as e:
        logger.exception(f"[{req_id}] ❌ 切换模型过程中发生严重错误")
        # 导入save_error_snapshot函数
        from .operations import save_error_snapshot
        await save_error_snapshot(f"model_switch_error_{req_id}")
        try:
            if original_prefs_str:
                logger.info(f"[{req_id}] 发生异常，尝试恢复 localStorage 至: {original_prompt_model or '未设置'}")
                await page.evaluate("(origPrefs) => localStorage.setItem('aiStudioUserPreference', origPrefs)", original_prefs_str)
                logger.info(f"[{req_id}] 异常恢复：导航到 '{new_chat_url}' 以应用恢复的 localStorage。")
                await page.goto(new_chat_url, wait_until="domcontentloaded", timeout=15000)
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(timeout=15000)
        except Exception as recovery_err:
            logger.error(f"[{req_id}] 异常后恢复 localStorage 失败: {recovery_err}")
        return False

def load_excluded_models(filename: str):
    """加载排除的模型列表"""
    import server
    excluded_model_ids = getattr(server, 'excluded_model_ids', set())
    
    excluded_file_path = os.path.join(os.path.dirname(__file__), '..', filename)
    try:
        if os.path.exists(excluded_file_path):
            with open(excluded_file_path, 'r', encoding='utf-8') as f:
                loaded_ids = {line.strip() for line in f if line.strip()}
            if loaded_ids:
                excluded_model_ids.update(loaded_ids)
                server.excluded_model_ids = excluded_model_ids
                logger.info(f"✅ 从 '{filename}' 加载了 {len(loaded_ids)} 个模型到排除列表: {excluded_model_ids}")
            else:
                logger.info(f"'{filename}' 文件为空或不包含有效的模型 ID，排除列表未更改。")
        else:
            logger.info(f"模型排除列表文件 '{filename}' 未找到，排除列表为空。")
    except Exception as e:
        logger.error(f"❌ 从 '{filename}' 加载排除模型列表时出错: {e}", exc_info=True)

async def _handle_initial_model_state_and_storage(page: AsyncPage):
    """处理初始模型状态和存储"""
    import server
    current_ai_studio_model_id = getattr(server, 'current_ai_studio_model_id', None)
    parsed_model_list = getattr(server, 'parsed_model_list', [])
    model_list_fetch_event = getattr(server, 'model_list_fetch_event', None)
    
    logger.info("--- (新) 处理初始模型状态, localStorage 和 isAdvancedOpen ---")
    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.aiStudioUserPreference 未找到。"
            logger.info(f"   判定需要刷新和存储更新: {reason_for_reload}")
        else:
            logger.info("   localStorage 中找到 'aiStudioUserPreference'。正在解析...")
            try:
                pref_obj = json.loads(initial_prefs_str)
                prompt_model_path = pref_obj.get("promptModel")
                is_advanced_open_in_storage = 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 = "localStorage.promptModel 无效或未设置。"
                    logger.info(f"   判定需要刷新和存储更新: {reason_for_reload}")
                else:
                    # 使用新的UI状态验证功能
                    ui_state = await _verify_ui_state_settings(page, "initial")
                    if ui_state['needsUpdate']:
                        needs_reload_and_storage_update = True
                        reason_for_reload = f"UI状态需要更新: isAdvancedOpen={ui_state['isAdvancedOpen']}, areToolsOpen={ui_state['areToolsOpen']} (期望: True)"
                        logger.info(f"   判定需要刷新和存储更新: {reason_for_reload}")
                    else:
                        server.current_ai_studio_model_id = prompt_model_path.split('/')[-1]
                        logger.info(f"   ✅ localStorage 有效且UI状态正确。初始模型 ID 从 localStorage 设置为: {server.current_ai_studio_model_id}")
            except json.JSONDecodeError:
                needs_reload_and_storage_update = True
                reason_for_reload = "解析 localStorage.aiStudioUserPreference JSON 失败。"
                logger.error(f"   判定需要刷新和存储更新: {reason_for_reload}")
        
        if needs_reload_and_storage_update:
            logger.info(f"   执行刷新和存储更新流程，原因: {reason_for_reload}")
            logger.info("   步骤 1: 调用 _set_model_from_page_display(set_storage=True) 更新 localStorage 和全局模型 ID...")
            await _set_model_from_page_display(page, set_storage=True)
            
            current_page_url = page.url
            logger.info(f"   步骤 2: 重新加载页面 ({current_page_url}) 以应用 isAdvancedOpen=true...")
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    logger.info(f"   尝试重新加载页面 (第 {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.info(f"   ✅ 页面已成功重新加载到: {page.url}")

                    # 页面重新加载后验证UI状态
                    logger.info(f"   页面重新加载完成，验证UI状态设置...")
                    reload_ui_state_success = await _verify_and_apply_ui_state(page, "reload")
                    if reload_ui_state_success:
                        logger.info(f"   ✅ 重新加载后UI状态验证成功")
                    else:
                        logger.warning(f"   ⚠️ 重新加载后UI状态验证失败")

                    break  # 成功则跳出循环
                except Exception as reload_err:
                    logger.warning(f"   ⚠️ 页面重新加载尝试 {attempt + 1}/{max_retries} 失败: {reload_err}")
                    if attempt < max_retries - 1:
                        logger.info(f"   将在5秒后重试...")
                        await asyncio.sleep(5)
                    else:
                        logger.error(f"   ❌ 页面重新加载在 {max_retries} 次尝试后最终失败: {reload_err}. 后续模型状态可能不准确。", exc_info=True)
                        from .operations import save_error_snapshot
                        await save_error_snapshot(f"initial_storage_reload_fail_attempt_{attempt+1}")
            
            logger.info("   步骤 3: 重新加载后，再次调用 _set_model_from_page_display(set_storage=False) 以同步全局模型 ID...")
            await _set_model_from_page_display(page, set_storage=False)
            logger.info(f"   ✅ 刷新和存储更新流程完成。最终全局模型 ID: {server.current_ai_studio_model_id}")
        else:
            logger.info("   localStorage 状态良好 (isAdvancedOpen=true, promptModel有效)，无需刷新页面。")
    except Exception as e:
        logger.error(f"❌ (新) 处理初始模型状态和 localStorage 时发生严重错误: {e}", exc_info=True)
        try:
            logger.warning("   由于发生错误，尝试回退仅从页面显示设置全局模型 ID (不写入localStorage)...")
            await _set_model_from_page_display(page, set_storage=False)
        except Exception as fallback_err:
            logger.error(f"   回退设置模型ID也失败: {fallback_err}")

async def _set_model_from_page_display(page: AsyncPage, set_storage: bool = False):
    """从页面显示设置模型"""
    import server
    current_ai_studio_model_id = getattr(server, 'current_ai_studio_model_id', None)
    parsed_model_list = getattr(server, 'parsed_model_list', [])
    model_list_fetch_event = getattr(server, 'model_list_fetch_event', None)
    
    try:
        logger.info("   尝试从页面显示元素读取当前模型名称...")
        model_name_locator = page.locator('[data-test-id="model-name"]')
        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.info(f"   页面当前显示模型名称 (原始: '{displayed_model_name_from_page_raw}', 清理后: '{displayed_model_name}')")
        
        found_model_id_from_display = None
        if model_list_fetch_event and not model_list_fetch_event.is_set():
            logger.info("   等待模型列表数据 (最多5秒) 以便转换显示名称...")
            try: 
                await asyncio.wait_for(model_list_fetch_event.wait(), timeout=5.0)
            except asyncio.TimeoutError: 
                logger.warning("   等待模型列表超时，可能无法准确转换显示名称为ID。")
        
        found_model_id_from_display = displayed_model_name
        logger.info(f"   页面显示的直接是模型ID: '{found_model_id_from_display}'")
        
        new_model_value = found_model_id_from_display
        if server.current_ai_studio_model_id != new_model_value:
            server.current_ai_studio_model_id = new_model_value
            logger.info(f"   全局 current_ai_studio_model_id 已更新为: {server.current_ai_studio_model_id}")
        else:
            logger.info(f"   全局 current_ai_studio_model_id ('{server.current_ai_studio_model_id}') 与从页面获取的值一致，未更改。")
        
        if set_storage:
            logger.info(f"   准备为页面状态设置 localStorage (确保 isAdvancedOpen=true)...")
            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("   解析现有 localStorage.aiStudioUserPreference 失败，将创建新的偏好设置。")
            
            # 使用新的强制设置功能
            logger.info(f"     应用强制UI状态设置...")
            ui_state_success = await _verify_and_apply_ui_state(page, "set_model")
            if not ui_state_success:
                logger.warning(f"     UI状态设置失败，使用传统方法")
                prefs_to_set["isAdvancedOpen"] = True
                prefs_to_set["areToolsOpen"] = True
            else:
                # 确保prefs_to_set也包含正确的设置
                prefs_to_set["isAdvancedOpen"] = True
                prefs_to_set["areToolsOpen"] = True
            logger.info(f"     强制 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
                logger.info(f"     设置 promptModel 为: {new_prompt_model_path} (基于找到的ID)")
            elif "promptModel" not in prefs_to_set:
                logger.warning(f"     无法从页面显示 '{displayed_model_name}' 找到模型ID，且 localStorage 中无现有 promptModel。promptModel 将不会被主动设置以避免潜在问题。")
            
            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.info(f"   ✅ localStorage.aiStudioUserPreference 已更新。isAdvancedOpen: {prefs_to_set.get('isAdvancedOpen')}, areToolsOpen: {prefs_to_set.get('areToolsOpen')} (期望: True), promptModel: '{prefs_to_set.get('promptModel', '未设置/保留原样')}'。")
    except Exception as e_set_disp:
        logger.error(f"   尝试从页面显示设置模型时出错: {e_set_disp}", exc_info=True)