"""浏览器 localStorage 协议的共享边界。""" import asyncio import json from dataclasses import dataclass from typing import Any from core.workspace import ( FAVORITES_STORAGE_KEY, HISTORY_STORAGE_KEY, LEGACY_STAGED_STORAGE_KEY, WORKSPACE_STORAGE_KEY, WorkspaceDataError, dump_workspace, empty_favorites, empty_history, merge_favorites, merge_history, merge_workspaces, migrate_legacy_workspace, normalize_favorites, normalize_history, normalize_workspace, ) from webui.constants import CONFIG_LS_KEY from webui.constants import ( LOCAL_STORAGE_MAX_READ_CHARS, LOCAL_STORAGE_READ_CHUNK_CHARS, LOCAL_STORAGE_RESTORE_CACHE, ) def storage_keys() -> dict[str, str]: return { 'workspace': WORKSPACE_STORAGE_KEY, 'legacy': LEGACY_STAGED_STORAGE_KEY, 'config': CONFIG_LS_KEY, 'history': HISTORY_STORAGE_KEY, 'favorites': FAVORITES_STORAGE_KEY, } def write_is_ready(name: str, states: dict[str, str], applying: set[str]) -> bool: """只有已安全恢复且未处于应用阶段的数据域允许写回浏览器。""" return states.get(name) == 'ready' and name not in applying async def backup_key(client, source_key: str, backup_key: str) -> bool: """在覆盖前仅创建一次浏览器内备份；已有备份绝不覆盖。""" source_key_js = json.dumps(source_key, ensure_ascii=False) backup_key_js = json.dumps(backup_key, ensure_ascii=False) try: status = await client.run_javascript( f'''(() => {{ const source = localStorage.getItem({source_key_js}); if (source === null) return 'missing'; if (localStorage.getItem({backup_key_js}) !== null) return 'exists'; try {{ localStorage.setItem({backup_key_js}, source); return 'created'; }} catch (error) {{ return `error:${{error && error.name ? error.name : 'unknown'}}`; }} }})()''', timeout=5.0, ) except Exception: return False return status in {'created', 'exists'} def storage_listener_script() -> str: keys = storage_keys() return f''' if (!window.__danbooruWorkspaceStorageListenerV1) {{ window.__danbooruWorkspaceStorageListenerV1 = true; window.addEventListener('storage', (event) => {{ const watchedKeys = new Set([{keys['workspace']!r}, {keys['history']!r}, {keys['favorites']!r}]); if (watchedKeys.has(event.key) && event.newValue !== event.oldValue) {{ if (window.confirm('工作区数据已在另一个标签页更新。是否重新加载当前页面以同步最新内容？')) window.location.reload(); }} }}); }} ''' async def read_prepared_value(client, connected, name: str, key: str, length) -> str | None: """按 UTF-16 安全边界分块读取准备好的浏览器缓存值。""" if length is None: return None if isinstance(length, bool) or not isinstance(length, (int, float)): raise RuntimeError(f'localStorage key {key!r} returned an invalid length') length = int(length) if length < 0 or length > LOCAL_STORAGE_MAX_READ_CHARS: raise WorkspaceDataError(f'localStorage key {key!r} exceeds the read limit') if length == 0: return '' name_js, key_js = json.dumps(name, ensure_ascii=False), json.dumps(key, ensure_ascii=False) cache_key_js = json.dumps(LOCAL_STORAGE_RESTORE_CACHE) chunks: list[str] = [] offset = 0 while offset < length: if not connected(): raise RuntimeError('client disconnected during localStorage restore') result = await client.run_javascript( f'''(() => {{ const cache = window[{cache_key_js}]; const value = cache && Object.prototype.hasOwnProperty.call(cache, {name_js}) ? cache[{name_js}] : localStorage.getItem({key_js}); if (value === null) return null; let end = Math.min(value.length, {offset + LOCAL_STORAGE_READ_CHUNK_CHARS}); if (end < value.length) {{ const last = value.charCodeAt(end - 1); if (last >= 0xD800 && last <= 0xDBFF) end += 1; }} return {{chunk: value.slice({offset}, end), next_offset: end}}; }})()''', timeout=5.0, ) if not isinstance(result, dict): raise RuntimeError(f'localStorage key {key!r} disappeared during restore') chunk, next_offset = result.get('chunk'), result.get('next_offset') if not isinstance(chunk, str) or not isinstance(next_offset, (int, float)): raise RuntimeError(f'localStorage key {key!r} returned an invalid chunk') next_offset = int(next_offset) if next_offset <= offset or next_offset > length: raise RuntimeError(f'localStorage key {key!r} returned an invalid offset') chunks.append(chunk) offset = next_offset return ''.join(chunks) def clear_restore_cache(client) -> None: try: client.run_javascript(f'delete window[{json.dumps(LOCAL_STORAGE_RESTORE_CACHE)}];') except RuntimeError: pass async def prepare_restore_snapshot(client, connected, keys: dict[str, str]) -> dict: """快照请求域；仅在浏览器内存中预压缩旧 history，绝不直接覆盖存储。""" if not connected(): raise RuntimeError('client is disconnected') keys_js = json.dumps(keys, ensure_ascii=False) cache_key_js = json.dumps(LOCAL_STORAGE_RESTORE_CACHE) result = await client.run_javascript( f'''(() => {{ const keys = {keys_js}, values = {{}}, manifest = {{}}; for (const [name, key] of Object.entries(keys)) {{ let value = localStorage.getItem(key), prepared = false; const originalLength = value === null ? null : value.length; if (name === 'history' && value) try {{ const data = JSON.parse(value); if (data && typeof data === 'object' && (data.schema_version === 1 || data.schema_version === 2) && Array.isArray(data.items)) {{ let changed = data.schema_version !== 2; const items = data.items.map((item) => {{ if (!item || typeof item !== 'object' || !item.workspace || typeof item.workspace !== 'object' || typeof item.query !== 'string' || !item.query.trim() || !item.settings || typeof item.settings !== 'object' || Array.isArray(item.settings)) return item; const query = item.query.trim().slice(0, 4000); const searchedAt = typeof item.searched_at === 'string' && item.searched_at ? item.searched_at : new Date().toISOString(); const oldQueries = item.workspace.queries; if (!Array.isArray(oldQueries) || oldQueries.length !== 1 || !oldQueries[0] || oldQueries[0].query !== query) changed = true; const workspace = {{...item.workspace, queries: [{{query, searched_at: searchedAt, settings: item.settings}}], updated_at: searchedAt}}; return {{...item, workspace_id: workspace.workspace_id, workspace}}; }}); if (changed) {{ value = JSON.stringify({{...data, schema_version: 2, items}}); prepared = true; }} }} }} catch (_) {{ /* Python validates and backs up corruption. */ }} values[name] = value; manifest[name] = {{length: value === null ? null : value.length, original_length: originalLength, prepared}}; }} window[{cache_key_js}] = values; return manifest; }})()''', timeout=5.0, ) if not isinstance(result, dict): raise RuntimeError('localStorage manifest is invalid') return result async def restore_staged_storage( controller: Any, names: tuple[str, ...], config_version: int, ) -> tuple[dict[str, str], set[str], list[str]]: """执行一次恢复尝试；控制器只提供页面状态和应用回调。""" unresolved = [ name for name in names if controller._storage_states.get(name) != 'ready' ] if not unresolved: return {}, set(), [] failures: dict[str, str] = {} persist: set[str] = set() warnings: list[str] = [] keys = controller._local_storage_keys() try: manifest = await controller._prepare_local_storage_restore(unresolved) except Exception as exc: message = str(exc) or type(exc).__name__ return {name: message for name in unresolved}, persist, warnings try: for name in unresolved: meta = manifest.get(name) if not isinstance(meta, dict): failures[name] = 'missing manifest entry' continue try: raw = await controller._read_local_storage_value( name, keys[name], meta.get('length') ) except Exception as exc: failures[name] = str(exc) or type(exc).__name__ continue controller._storage_raw_values[name] = raw finally: controller._clear_local_storage_restore_cache() if 'legacy' in unresolved and 'legacy' not in failures: controller._storage_states['legacy'] = 'ready' if 'config' in unresolved and 'config' not in failures: raw_config = controller._storage_raw_values.get('config') config_dirty = 'config' in controller._storage_session_dirty try: cfg = json.loads(raw_config) if raw_config else {} if not isinstance(cfg, dict): raise WorkspaceDataError('config must be a JSON object') except Exception as exc: config_key = keys['config'] if raw_config and not await controller._backup_local_storage_key( config_key, f'{config_key}_corrupt_backup' ): failures['config'] = f'corrupt config backup failed: {exc}' else: warnings.append('config_corrupt') persist.add('config') else: if cfg and cfg.get('version') != config_version: warnings.append('config_schema_migrated') persist.add('config') if not config_dirty: controller._storage_applying.add('config') try: controller._apply_config_state(cfg) finally: controller._storage_applying.discard('config') else: persist.add('config') if 'config' not in failures: controller._storage_states['config'] = 'ready' if 'workspace' in unresolved and 'workspace' not in failures: raw_workspace = controller._storage_raw_values.get('workspace') if not raw_workspace and any( controller._storage_states.get(name) != 'ready' for name in ('legacy', 'config') ): failures['workspace'] = 'legacy workspace inputs are not available' else: workspace_warnings: list[str] = [] try: if raw_workspace: workspace, workspace_warnings = normalize_workspace(raw_workspace) else: workspace, workspace_warnings = migrate_legacy_workspace( controller._storage_raw_values.get('legacy'), controller._storage_raw_values.get('config'), ) persist.add('workspace') except WorkspaceDataError as exc: backed_up = await controller._backup_local_storage_key( keys['workspace'], f"{keys['workspace']}_corrupt_backup" ) if raw_workspace and not backed_up: failures['workspace'] = f'corrupt workspace backup failed: {exc}' else: workspace, migration_warnings = migrate_legacy_workspace( controller._storage_raw_values.get('legacy'), controller._storage_raw_values.get('config'), ) workspace_warnings = ['workspace_corrupt'] + migration_warnings persist.add('workspace') if 'workspace' not in failures: if 'workspace' in controller._storage_session_dirty: workspace = merge_workspaces( controller.workspace_state, workspace, origin='local_restore', source='浏览器本地恢复', ) persist.add('workspace') controller._storage_applying.add('workspace') try: controller._apply_workspace_state( workspace, persist=False, refresh_recommendations=False ) finally: controller._storage_applying.discard('workspace') warnings.extend(workspace_warnings) if workspace_warnings: persist.add('workspace') controller._storage_states['workspace'] = 'ready' if raw_workspace: controller._storage_states['legacy'] = 'ready' failures.pop('legacy', None) if 'history' in unresolved and 'history' not in failures: raw_history = controller._storage_raw_values.get('history') history_prepared = bool( isinstance(manifest.get('history'), dict) and manifest['history'].get('prepared') ) try: history, history_warnings = normalize_history(raw_history) except WorkspaceDataError as exc: backed_up = await controller._backup_local_storage_key( keys['history'], f"{keys['history']}_corrupt_backup" ) if raw_history and not backed_up: failures['history'] = f'corrupt history backup failed: {exc}' else: history, history_warnings = empty_history(), ['history_corrupt'] persist.add('history') if 'history' not in failures: if history_prepared: if not await controller._backup_history_before_compaction(): failures['history'] = 'legacy history backup failed' else: history_warnings.extend([ 'history_schema_migrated', 'history_workspace_queries_compacted', ]) persist.add('history') if 'history' not in failures: if 'history' in controller._storage_session_dirty: history = merge_history(controller.search_history, history) persist.add('history') controller.search_history = history warnings.extend(history_warnings) if history_warnings: persist.add('history') controller._storage_states['history'] = 'ready' if 'favorites' in unresolved and 'favorites' not in failures: raw_favorites = controller._storage_raw_values.get('favorites') try: favorites, favorite_warnings = normalize_favorites(raw_favorites) except WorkspaceDataError as exc: backed_up = await controller._backup_local_storage_key( keys['favorites'], f"{keys['favorites']}_corrupt_backup" ) if raw_favorites and not backed_up: failures['favorites'] = f'corrupt favorites backup failed: {exc}' else: favorites, favorite_warnings = empty_favorites(), ['favorites_corrupt'] persist.add('favorites') if 'favorites' not in failures: if 'favorites' in controller._storage_session_dirty: favorites = merge_favorites(controller.favorites, favorites) persist.add('favorites') controller.favorites = favorites warnings.extend(favorite_warnings) if favorite_warnings: persist.add('favorites') controller._storage_states['favorites'] = 'ready' for name in failures: controller._storage_states[name] = 'failed' controller._update_undo_buttons() controller._update_workspace_counts() return failures, persist, warnings @dataclass(frozen=True) class RestoreRetryResult: """一次恢复任务的无界面结果，交由页面层决定日志和提示。""" completed: bool client_stopped: bool failures: dict[str, str] warnings: list[str] async def restore_with_retries( controller: Any, retry_delays: tuple[float, ...], ) -> RestoreRetryResult: """在连接可用时执行有限次恢复，不触碰 NiceGUI 的通知上下文。""" last_failures: dict[str, str] = {} all_warnings: list[str] = [] for delay in retry_delays: if delay: await asyncio.sleep(delay) if not controller._client_alive(): return RestoreRetryResult(False, True, last_failures, all_warnings) if not controller._client_connected(): last_failures = {'connection': 'client is disconnected'} continue controller._storage_restoring = True try: failures, persist, warnings = await controller._restore_staged_tags() finally: controller._storage_restoring = False persist_restored_storage(controller, persist) all_warnings.extend(warnings) last_failures = failures if not failures: return RestoreRetryResult(True, False, {}, all_warnings) return RestoreRetryResult(False, False, last_failures, all_warnings) def persist_restored_storage(controller: Any, names: set[str]) -> None: """按既定顺序写回已恢复的数据域，具体写入仍由页面控制器完成。""" for name, saver_name in ( ('config', '_save_config'), ('workspace', '_save_staged_tags'), ('history', '_save_history'), ('favorites', '_save_favorites'), ): if name in names: getattr(controller, saver_name)() def flush_storage_session_changes(controller: Any) -> None: """只刷新已安全恢复的数据域，失败或未恢复域绝不覆盖浏览器旧值。""" ready_dirty = { name for name in controller._storage_session_dirty if controller._storage_states.get(name) == 'ready' } persist_restored_storage(controller, ready_dirty) def start_storage_restore_task(controller: Any, worker) -> Any: """仅在仍有待恢复域或待写回数据时创建一个恢复任务。""" task = controller._storage_restore_task if task is not None and not task.done(): return task needs_work = any( state != 'ready' for state in controller._storage_states.values() ) or bool(controller._storage_session_dirty) if not needs_work: return None controller._storage_restore_task = asyncio.create_task(worker()) return controller._storage_restore_task def pause_storage_restore(controller: Any) -> None: """取消当前恢复任务；重连后的重启决定仍由页面生命周期处理。""" task = controller._storage_restore_task if task is not None and not task.done(): task.cancel() def finish_storage_restore_task( controller: Any, current_task: Any, was_cancelled: bool, ) -> bool: """清理当前任务，并告知页面层是否应在已连接时重新启动恢复。""" controller._storage_restoring = False if controller._storage_restore_task is not current_task: return False controller._storage_restore_task = None return was_cancelled and controller._client_connected() def schedule_workspace_persist(controller: Any, debounce_seconds: float) -> None: """去抖写入工作区；恢复未就绪或写入失败时保留原有浏览器数据。""" if not controller._storage_write_allowed('workspace'): return if controller._workspace_save_task and not controller._workspace_save_task.done(): controller._workspace_save_task.cancel() async def persist() -> None: try: await asyncio.sleep(debounce_seconds) except asyncio.CancelledError: return try: data = dump_workspace(controller.workspace_state) except WorkspaceDataError as exc: print(f'[UI] 工作区保存前校验失败: {exc}', flush=True) return if not controller._storage_write_allowed('workspace'): return try: controller.client.run_javascript( f"localStorage.setItem({json.dumps(WORKSPACE_STORAGE_KEY)}, {json.dumps(data)});" ) except RuntimeError: controller._storage_session_dirty.add('workspace') return controller._storage_session_dirty.discard('workspace') controller._workspace_save_task = asyncio.ensure_future(persist())