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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 等。"""
|
||||
|
||||
Reference in New Issue
Block a user