feat: persist browser sessions and update admin workflows
This commit is contained in:
@@ -34,6 +34,7 @@ class BrowserSession:
|
||||
lock: asyncio.Lock
|
||||
cdp_session: Any = None
|
||||
captured_headers: list[dict] = None # auth headers from CDP
|
||||
last_saved_state_at: float = 0.0
|
||||
|
||||
|
||||
class BrowserSessionService:
|
||||
@@ -92,12 +93,15 @@ class BrowserSessionService:
|
||||
self._profiles.pop(profile_key, None)
|
||||
# Idle cleanup: close stale sessions before spawning new ones
|
||||
await self._evict_idle_sessions()
|
||||
|
||||
context = await self._playwright.chromium.launch_persistent_context(
|
||||
str(self._profile_dir(profile_key)),
|
||||
headless=get_settings().browser_headless,
|
||||
viewport={"width": width, "height": height},
|
||||
color_scheme="dark",
|
||||
args=["--no-sandbox", "--disable-dev-shm-usage"],
|
||||
)
|
||||
await self._restore_session_state(context, profile_key)
|
||||
# Grant clipboard access for the page origin
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
@@ -137,6 +141,11 @@ class BrowserSessionService:
|
||||
self._touch(session_id)
|
||||
async with session.lock:
|
||||
self._ensure_open(session)
|
||||
if session.profile_key and not session.profile_key.startswith("auth-capture-"):
|
||||
now = time.monotonic()
|
||||
if now - session.last_saved_state_at > 10.0:
|
||||
await self._save_session_state(session)
|
||||
session.last_saved_state_at = now
|
||||
return await session.page.screenshot(type="jpeg", quality=65, full_page=False)
|
||||
|
||||
async def event(
|
||||
@@ -188,6 +197,12 @@ class BrowserSessionService:
|
||||
await page.set_viewport_size({"width": width, "height": height})
|
||||
else:
|
||||
raise ValueError("Unsupported browser event")
|
||||
if session.profile_key and not session.profile_key.startswith("auth-capture-"):
|
||||
now = time.monotonic()
|
||||
if now - session.last_saved_state_at > 5.0:
|
||||
await self._save_session_state(session)
|
||||
session.last_saved_state_at = now
|
||||
|
||||
if not include_state:
|
||||
return None
|
||||
return await self._session_state(session)
|
||||
@@ -242,6 +257,15 @@ class BrowserSessionService:
|
||||
session = self._discard_session(session_id)
|
||||
if not session:
|
||||
return
|
||||
|
||||
# 在完全关闭 context 前,强制将最新的状态落盘保存
|
||||
if session.profile_key and not session.profile_key.startswith("auth-capture-"):
|
||||
try:
|
||||
if not session.page.is_closed():
|
||||
await self._save_session_state(session)
|
||||
except Exception as exc:
|
||||
logger.debug("failed to save state during close: %s", exc)
|
||||
|
||||
# Detach CDP session if active
|
||||
if session.cdp_session:
|
||||
try:
|
||||
@@ -261,6 +285,7 @@ class BrowserSessionService:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
# Cancel the background eviction loop
|
||||
if self._evict_task is not None and not self._evict_task.done():
|
||||
@@ -524,6 +549,9 @@ class BrowserSessionService:
|
||||
profile.mkdir(parents=True, exist_ok=True)
|
||||
return profile
|
||||
|
||||
def _cookies_path(self, profile_key: str) -> Path:
|
||||
return self._profile_dir(profile_key) / "session-cookies.json"
|
||||
|
||||
def _profile_key(self, custom_page_id: int, url: str) -> str:
|
||||
parsed = urlparse(url)
|
||||
origin = f"{parsed.scheme}-{parsed.netloc}".lower()
|
||||
@@ -553,6 +581,7 @@ class BrowserSessionService:
|
||||
str(self._profile_dir(profile_key)),
|
||||
headless=get_settings().browser_headless,
|
||||
viewport={"width": width, "height": height},
|
||||
color_scheme="dark",
|
||||
args=["--no-sandbox", "--disable-dev-shm-usage"],
|
||||
)
|
||||
# Grant clipboard access for the page origin
|
||||
@@ -613,5 +642,85 @@ class BrowserSessionService:
|
||||
except Exception as exc:
|
||||
logger.debug("CDP capture not available: %s", exc)
|
||||
|
||||
async def _save_session_state(self, session: BrowserSession) -> None:
|
||||
if not session.profile_key or session.profile_key.startswith("auth-capture-"):
|
||||
return
|
||||
try:
|
||||
state = await session.context.storage_state()
|
||||
cookies_path = self._cookies_path(session.profile_key)
|
||||
import json
|
||||
import tempfile
|
||||
import os
|
||||
# Ensure parent directories exist
|
||||
cookies_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temp_fd, temp_path = tempfile.mkstemp(dir=str(cookies_path.parent))
|
||||
try:
|
||||
with os.fdopen(temp_fd, 'w', encoding='utf-8') as f:
|
||||
json.dump(state, f, ensure_ascii=False, indent=2)
|
||||
os.replace(temp_path, cookies_path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.debug("failed to save session state for %s: %s", session.profile_key, exc)
|
||||
|
||||
async def _restore_session_state(self, context: Any, profile_key: str) -> None:
|
||||
if profile_key.startswith("auth-capture-"):
|
||||
return
|
||||
cookies_path = self._cookies_path(profile_key)
|
||||
if not cookies_path.exists() or cookies_path.stat().st_size == 0:
|
||||
return
|
||||
try:
|
||||
import json
|
||||
import time
|
||||
with open(cookies_path, 'r', encoding='utf-8') as f:
|
||||
state = json.load(f)
|
||||
cookies = state.get("cookies", [])
|
||||
if cookies:
|
||||
now = time.time()
|
||||
valid_cookies = []
|
||||
for c in cookies:
|
||||
expires = c.get("expires")
|
||||
if expires is not None and expires > 0 and expires <= now:
|
||||
continue
|
||||
if expires is not None and expires <= 0:
|
||||
c.pop("expires", None)
|
||||
valid_cookies.append(c)
|
||||
if valid_cookies:
|
||||
await context.add_cookies(valid_cookies)
|
||||
logger.info("restored %d cookies for profile %s", len(valid_cookies), profile_key)
|
||||
|
||||
# 还原 LocalStorage
|
||||
origins = state.get("origins", [])
|
||||
if origins:
|
||||
origins_json = json.dumps(origins)
|
||||
init_script = f"""
|
||||
(() => {{
|
||||
try {{
|
||||
const origins = {origins_json};
|
||||
const currentOrigin = window.location.origin;
|
||||
const target = origins.find(o => o.origin === currentOrigin);
|
||||
if (target && target.localStorage) {{
|
||||
for (const item of target.localStorage) {{
|
||||
try {{
|
||||
window.localStorage.setItem(item.name, item.value);
|
||||
}} catch (e) {{
|
||||
console.error('Failed to restore localStorage key', item.name, e);
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
}} catch (err) {{
|
||||
console.error('LocalStorage restore initialization script failed', err);
|
||||
}}
|
||||
}})();
|
||||
"""
|
||||
await context.add_init_script(init_script)
|
||||
logger.info("registered LocalStorage init script for profile %s (origins: %d)", profile_key, len(origins))
|
||||
except Exception as exc:
|
||||
logger.warning("failed to restore cookies/state for profile %s: %s", profile_key, exc)
|
||||
|
||||
|
||||
browser_sessions = BrowserSessionService()
|
||||
|
||||
Reference in New Issue
Block a user