Spaces:
Paused
Paused
File size: 2,119 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 | # --- browser_utils/initialization/auth.py ---
"""
Authentication Saving Module - Simplified Version
Handles saving authentication state after login. Automatically saves to SAVED_AUTH_DIR.
"""
import asyncio
import logging
import os
import time
from config import SAVED_AUTH_DIR
logger = logging.getLogger("AIStudioProxyServer")
async def wait_for_model_list_and_handle_auth_save(temp_context, launch_mode, loop):
"""Wait for model list response and handle authentication saving"""
from api_utils.server_state import state
# Wait for model list response to confirm login success
logger.info("Waiting for model list response to confirm login success...")
try:
await asyncio.wait_for(state.model_list_fetch_event.wait(), timeout=30.0)
logger.info("Model list response detected, login confirmed!")
except asyncio.TimeoutError:
logger.warning(
"Timeout waiting for model list response, but continuing with auth save..."
)
# Determine filename: env var > auto-generate
filename = os.environ.get("SAVE_AUTH_FILENAME", "").strip()
if not filename:
filename = f"auth_auto_{int(time.time())}"
await _save_auth_state(temp_context, filename)
async def _save_auth_state(temp_context, filename: str):
"""Unified authentication saving function"""
os.makedirs(SAVED_AUTH_DIR, exist_ok=True)
if not filename.endswith(".json"):
filename += ".json"
auth_save_path = os.path.join(SAVED_AUTH_DIR, filename)
print("\n" + "=" * 50, flush=True)
print("Login successful! Saving authentication state...", flush=True)
try:
await temp_context.storage_state(path=auth_save_path)
logger.info(f"Authentication state saved to: {auth_save_path}")
print(f"Authentication state saved to: {auth_save_path}", flush=True)
except asyncio.CancelledError:
raise
except Exception as e:
logger.error(f"Failed to save authentication state: {e}", exc_info=True)
print(f"Failed to save authentication state: {e}", flush=True)
print("=" * 50 + "\n", flush=True)
|