feat: persist browser sessions and update admin workflows

This commit is contained in:
liumangmang
2026-05-29 16:00:43 +08:00
parent e3151a7ea6
commit c5778bb3e7
19 changed files with 829 additions and 369 deletions
@@ -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()
+31
View File
@@ -46,6 +46,7 @@ def _check_upstream(upstream_id: int) -> None:
auth_config = json.loads(upstream.auth_config_json or "{}")
was_unhealthy = upstream.last_status == "unhealthy"
balance_alert_triggered = False
snapshot = None
changes = None
@@ -76,6 +77,14 @@ def _check_upstream(upstream_id: int) -> None:
if balance is not None:
upstream.balance = balance
upstream.balance_updated_at = datetime.now(timezone.utc)
# ── 余额告警阈值检查 ──
threshold = upstream.balance_alert_threshold
if threshold is not None and threshold > 0:
if balance < threshold and not upstream.balance_alert_notified:
upstream.balance_alert_notified = True
balance_alert_triggered = True
elif balance >= threshold and upstream.balance_alert_notified:
upstream.balance_alert_notified = False
except Exception as exc:
# failure path
upstream.consecutive_failures = (upstream.consecutive_failures or 0) + 1
@@ -152,6 +161,12 @@ def _check_upstream(upstream_id: int) -> None:
_notify_rate_changed(upstream_id, upstream.name, upstream.base_url, changes)
_sync_website_bindings(upstream_id, changes)
if balance_alert_triggered:
_notify_balance_low(
upstream_id, upstream.name, upstream.base_url,
upstream.balance, upstream.balance_alert_threshold,
)
def _notify_status(
upstream_id: int,
@@ -184,6 +199,22 @@ def _notify_rate_changed(
db.close()
def _notify_balance_low(
upstream_id: int,
upstream_name: str,
base_url: str,
balance: float,
threshold: float,
) -> None:
db = SessionLocal()
try:
webhook_service.send_balance_low(db, upstream_id, upstream_name, base_url, balance, threshold)
except Exception:
logger.exception("balance low webhook failed for upstream %s", upstream_name)
finally:
db.close()
def _sync_upstream_keys(upstream_id: int, snapshot: dict[str, Any], captured_at: datetime) -> None:
"""上游检测成功后同步 SmartUp Key 状态(远端删除/分组删除)。"""
db = SessionLocal()
+38
View File
@@ -15,6 +15,7 @@ from app.utils.dingtalk import (
format_dingtalk_rate_changed,
format_dingtalk_website_rate_changed,
format_dingtalk_status,
format_dingtalk_balance_low,
)
@@ -185,6 +186,43 @@ def send_status_event(
_log(db, wh, event, generic_payload, "failed", str(exc))
def send_balance_low(
db: Session,
upstream_id: int,
upstream_name: str,
base_url: str,
balance: float,
threshold: float,
) -> None:
webhooks = (
db.query(WebhookConfig)
.filter(WebhookConfig.enabled == True)
.all()
)
event = "upstream_balance_low"
changed_at = _now_iso()
generic_payload = {
"event": event,
"upstream": {"id": upstream_id, "name": upstream_name, "base_url": base_url},
"balance": balance,
"threshold": threshold,
"changed_at": changed_at,
}
for wh in webhooks:
events = json.loads(wh.events_json or "[]")
if event not in events:
continue
try:
if wh.type == "dingtalk":
msg = format_dingtalk_balance_low(upstream_name, balance, threshold, changed_at)
resp_text = _send_dingtalk(wh.url, wh.secret, msg)
else:
resp_text = _send_generic(wh.url, generic_payload)
_log(db, wh, event, generic_payload, "success", resp_text)
except Exception as exc:
_log(db, wh, event, generic_payload, "failed", str(exc))
def send_test_notification(db: Session, webhook: WebhookConfig) -> tuple[bool, str]:
payload = {
"event": "test",
+6
View File
@@ -223,6 +223,12 @@ class Sub2ApiWebsiteClient:
data = _unwrap_data(resp)
return data if isinstance(data, dict) else {"value": data}
def update_account(self, account_id: str, body: dict[str, Any], endpoint: str = "/accounts") -> dict[str, Any]:
"""更新远端账号(仅传入需要变更的字段)。"""
resp = self._request("PUT", f"{endpoint}/{account_id}", body)
data = _unwrap_data(resp)
return data if isinstance(data, dict) else {"value": data}
@staticmethod
def _unwrap_list(value: dict) -> list | None:
"""递归展开嵌套的列表包装:data.items、data.data、items、accounts 等。"""