feat: remote browser login persistence + balance display + UI consistency

- Retain login state in remote browser profiles (don't delete on disconnect)
- Add GET /api/browser-sessions/{id}/clipboard for clipboard sync
- Add POST /api/browser-sessions/{id}/autofill-login for manual credential fill
- Add DELETE /api/browser-sessions/profiles/{custom_page_id} for login clear
- Add balance tracking with configurable divisor (balance_divisor)
- Health check on session reuse, idle TTL eviction, background cleanup
- Add first-frame watchdog (10s timeout) to prevent infinite loading
- Reconnect browser on active=true when session was closed
- UI: uniform text-only inline buttons (websites + upstreams pages)
- Fix page switch race with closingRemoteSessionPromise
This commit is contained in:
liumangmang
2026-05-20 09:44:20 +08:00
parent 4c71148ff9
commit 6cc797f915
16 changed files with 773 additions and 52 deletions
+20
View File
@@ -29,6 +29,7 @@ def init_db():
from app.models import admin_user, upstream, snapshot, webhook_config, notification_log, custom_page, website, revoked_token # noqa: F401
Base.metadata.create_all(bind=engine)
_migrate_custom_pages()
_migrate_upstreams()
def _migrate_custom_pages():
@@ -67,3 +68,22 @@ def _migrate_custom_pages():
if "linked_upstream_id" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN linked_upstream_id INTEGER"))
def _migrate_upstreams():
"""Apply SQLite-safe migrations to the upstreams table."""
inspector = inspect(engine)
if "upstreams" not in inspector.get_table_names():
return
columns = {col["name"] for col in inspector.get_columns("upstreams")}
with engine.begin() as conn:
if "balance" not in columns:
conn.execute(text("ALTER TABLE upstreams ADD COLUMN balance FLOAT"))
if "balance_updated_at" not in columns:
conn.execute(text("ALTER TABLE upstreams ADD COLUMN balance_updated_at DATETIME"))
if "balance_endpoint" not in columns:
conn.execute(text("ALTER TABLE upstreams ADD COLUMN balance_endpoint VARCHAR(256) NOT NULL DEFAULT ''"))
if "balance_response_path" not in columns:
conn.execute(text("ALTER TABLE upstreams ADD COLUMN balance_response_path VARCHAR(256) NOT NULL DEFAULT ''"))
if "balance_divisor" not in columns:
conn.execute(text("ALTER TABLE upstreams ADD COLUMN balance_divisor FLOAT NOT NULL DEFAULT 1.0"))
+7 -1
View File
@@ -1,6 +1,6 @@
from datetime import datetime, timezone
from typing import Optional
from sqlalchemy import Integer, String, Boolean, DateTime, Text
from sqlalchemy import Integer, String, Boolean, DateTime, Text, Float
from sqlalchemy.orm import mapped_column, Mapped
from app.database import Base
@@ -26,6 +26,12 @@ class Upstream(Base):
last_checked_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
last_error: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
consecutive_failures: Mapped[int] = mapped_column(Integer, default=0)
# Balance tracking
balance: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
balance_updated_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
balance_endpoint: Mapped[str] = mapped_column(String(256), default="")
balance_response_path: Mapped[str] = mapped_column(String(256), default="")
balance_divisor: Mapped[float] = mapped_column(Float, default=1.0)
created_at: Mapped[datetime] = mapped_column(DateTime, default=lambda: datetime.now(timezone.utc))
updated_at: Mapped[datetime] = mapped_column(
DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)
+102
View File
@@ -131,11 +131,113 @@ async def get_selection(session_id: str, _=Depends(get_current_user)):
raise _error_from_browser(exc)
class BrowserClipboardResponse(BaseModel):
text: Optional[str] = None
error: Optional[str] = None
@router.get("/{session_id}/clipboard", response_model=BrowserClipboardResponse)
async def session_clipboard(session_id: str, _=Depends(get_current_user)):
"""Read text from the remote browser's clipboard."""
from fastapi.responses import JSONResponse
try:
text, error = await browser_sessions.read_clipboard(session_id)
body: dict[str, Any] = {}
if text:
body["text"] = text
elif error == "denied":
body["error"] = "远程浏览器未授予剪贴板读取权限"
elif error == "read_failed":
body["error"] = "读取远程剪贴板时发生内部错误"
else:
if error:
logger.warning("clipboard read error for %s: %s", session_id[:12], error)
body["error"] = "远程剪贴板为空"
return JSONResponse(content=body, headers={"Cache-Control": "no-store"})
except Exception as exc:
raise _error_from_browser(exc)
class AutofillLoginResponse(BaseModel):
success: bool
message: str
@router.post("/{session_id}/autofill-login", response_model=AutofillLoginResponse)
async def autofill_login(session_id: str, _=Depends(get_current_user)):
"""Manually trigger login autofill for the remote browser session.
Uses the linked custom page's saved credentials. Never returns passwords.
"""
try:
session_state = await browser_sessions.state(session_id)
except Exception as exc:
raise _error_from_browser(exc)
from app.database import SessionLocal as _Db
from app.models.custom_page import CustomPage
db = _Db()
try:
page = db.query(CustomPage).filter(
CustomPage.id == session_state["custom_page_id"]
).first()
if not page or not page.enabled:
raise HTTPException(400, "linked custom page is not available")
if page.access_mode != "remote_browser":
raise HTTPException(400, "linked custom page is not in remote browser mode")
if not page.login_autofill_enabled:
return AutofillLoginResponse(success=False, message="该页面未启用自动填充登录")
if not page.login_username or not page.login_password:
return AutofillLoginResponse(success=False, message="该页面未保存账号密码")
login_config = {
"enabled": True,
"username": page.login_username,
"password": page.login_password,
"username_selector": page.login_username_selector,
"password_selector": page.login_password_selector,
"submit_selector": page.login_submit_selector,
}
filled = await browser_sessions.autofill_login(session_id, login_config)
if filled:
return AutofillLoginResponse(success=True, message="已填入账号密码")
return AutofillLoginResponse(
success=False,
message="未找到登录输入框,请先关闭弹窗或进入登录页后重试",
)
finally:
db.close()
@router.delete("/{session_id}", status_code=204)
async def close_session(session_id: str, _=Depends(get_current_user)):
await browser_sessions.close(session_id)
@router.delete("/profiles/{custom_page_id}", status_code=204)
async def clear_profile(custom_page_id: int, _=Depends(get_current_user)):
"""Close active session for the page and delete its profile directory.
On next open the browser starts fresh, losing login state.
"""
from app.models.custom_page import CustomPage
from app.database import SessionLocal as _Db
db = _Db()
try:
page = db.query(CustomPage).filter(CustomPage.id == custom_page_id).first()
if not page or not page.enabled:
raise HTTPException(404, "custom page not found")
if page.access_mode != "remote_browser":
raise HTTPException(400, "custom page is not in remote browser mode")
try:
await browser_sessions.clear_profile(custom_page_id, page.url)
except RuntimeError as exc:
raise HTTPException(500, str(exc))
finally:
db.close()
# ——— WebSocket stream ———
# Frame interval & diff detection
_WS_MIN_INTERVAL = 0.10
+14 -1
View File
@@ -270,7 +270,20 @@ async def refresh_auth(pid: int, db: Session = Depends(get_db), _=Depends(get_cu
existing_config["new_api_user"] = candidate["new_api_user"]
elif ctype == "bearer_token":
upstream.auth_type = "bearer"
existing_config["token"] = candidate.get("value", "")
raw = candidate.get("value", "")
# Clean up: strip whitespace, remove "Bearer " prefix if present
token = raw.strip()
if token.startswith("Bearer "):
token = token[7:].strip()
# Validate token can be used as HTTP header value
try:
token.encode("latin-1")
except UnicodeEncodeError:
return RefreshAuthResponse(
success=False,
message=f"提取到的 Token 含有非 HTTP 标头字符,请确认已在远程浏览器中正确登录并重试",
)
existing_config["token"] = token
elif ctype == "api_key":
upstream.auth_type = "api_key"
existing_config["key"] = candidate.get("value", "")
+31
View File
@@ -2,9 +2,12 @@
from __future__ import annotations
import json
import logging
from datetime import datetime, timezone
from typing import List
logger = logging.getLogger(__name__)
from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session
@@ -55,6 +58,11 @@ def _to_response(u: Upstream) -> UpstreamResponse:
last_status=u.last_status,
last_checked_at=u.last_checked_at,
last_error=u.last_error,
balance=u.balance,
balance_updated_at=u.balance_updated_at,
balance_endpoint=u.balance_endpoint or "",
balance_response_path=u.balance_response_path or "",
balance_divisor=u.balance_divisor or 1.0,
created_at=u.created_at,
updated_at=u.updated_at,
)
@@ -82,6 +90,9 @@ def create_upstream(
enabled=body.enabled,
check_interval_seconds=body.check_interval_seconds,
timeout_seconds=body.timeout_seconds,
balance_endpoint=body.balance_endpoint,
balance_response_path=body.balance_response_path,
balance_divisor=body.balance_divisor,
)
db.add(u)
db.commit()
@@ -156,6 +167,16 @@ def test_upstream(uid: int, db: Session = Depends(get_db), _=Depends(get_current
try:
client.login()
groups = client.get_available_groups(u.groups_endpoint)
# Also try balance if configured
if u.balance_endpoint and u.balance_response_path:
try:
raw_balance = client.get_balance(u.balance_endpoint, u.balance_response_path)
if raw_balance is not None:
divisor = u.balance_divisor or 1.0
u.balance = raw_balance / divisor
u.balance_updated_at = datetime.now(timezone.utc) if raw_balance is not None else None
except Exception as exc:
logger.warning("upstream %s balance fetch failed during test: %s", u.name, exc)
u.last_status = "healthy"
u.last_error = None
u.last_checked_at = datetime.now(timezone.utc)
@@ -189,6 +210,16 @@ def check_now(uid: int, db: Session = Depends(get_db), _=Depends(get_current_use
groups = client.get_available_groups(u.groups_endpoint)
raw_rates = client.get_group_rates(u.rate_endpoint)
snapshot = build_snapshot(u.id, u.base_url, u.api_prefix, groups, raw_rates)
# Also try balance if configured
if u.balance_endpoint and u.balance_response_path:
try:
raw_balance = client.get_balance(u.balance_endpoint, u.balance_response_path)
if raw_balance is not None:
divisor = u.balance_divisor or 1.0
u.balance = raw_balance / divisor
u.balance_updated_at = datetime.now(timezone.utc) if raw_balance is not None else None
except Exception as exc:
logger.warning("upstream %s balance fetch failed during check-now: %s", u.name, exc)
except Exception as exc:
u.consecutive_failures = (u.consecutive_failures or 0) + 1
u.last_error = str(exc)
+11
View File
@@ -29,6 +29,9 @@ class UpstreamCreate(BaseModel):
enabled: bool = True
check_interval_seconds: int = 600
timeout_seconds: int = 30
balance_endpoint: str = ""
balance_response_path: str = ""
balance_divisor: float = 1.0
class UpstreamUpdate(BaseModel):
@@ -42,6 +45,9 @@ class UpstreamUpdate(BaseModel):
enabled: Optional[bool] = None
check_interval_seconds: Optional[int] = None
timeout_seconds: Optional[int] = None
balance_endpoint: Optional[str] = None
balance_response_path: Optional[str] = None
balance_divisor: Optional[float] = None
class UpstreamResponse(BaseModel):
@@ -59,6 +65,11 @@ class UpstreamResponse(BaseModel):
last_status: str
last_checked_at: Optional[datetime]
last_error: Optional[str]
balance: Optional[float] = None
balance_updated_at: Optional[datetime] = None
balance_endpoint: str = ""
balance_response_path: str = ""
balance_divisor: float = 1.0
created_at: datetime
updated_at: datetime
+198 -18
View File
@@ -37,11 +37,18 @@ class BrowserSession:
class BrowserSessionService:
# Idle TTL: close sessions that haven't had activity for this long
IDLE_TTL_SECONDS = 1800 # 30 minutes
# Cap: max concurrent persistent sessions (excludes auth-capture)
MAX_SESSIONS = 10
def __init__(self) -> None:
self._playwright: Optional[Any] = None
self._sessions: dict[str, BrowserSession] = {}
self._profiles: dict[str, str] = {}
self._lock = asyncio.Lock()
self._last_event_at: dict[str, float] = {}
self._evict_task: Optional[asyncio.Task[None]] = None
async def create(
self,
@@ -61,21 +68,43 @@ class BrowserSessionService:
existing_id = self._profiles.get(profile_key)
existing = self._sessions.get(existing_id or "")
if existing and not existing.page.is_closed():
async with existing.lock:
await existing.page.set_viewport_size({"width": width, "height": height})
if existing.page.url == "about:blank":
await existing.page.goto(url, wait_until="domcontentloaded", timeout=45000)
await self._autofill_login(existing.page, login_config)
await self._reset_page_zoom(existing)
# Health check: verify session can actually serve content
healthy = True
try:
async with existing.lock:
url_before = existing.page.url
await existing.page.evaluate("1") # ping
await existing.page.screenshot(type="jpeg", quality=10, timeout=5000)
await existing.page.set_viewport_size({"width": width, "height": height})
if url_before == "about:blank":
await existing.page.goto(url, wait_until="domcontentloaded", timeout=45000)
await self._autofill_login(existing.page, login_config)
await self._reset_page_zoom(existing)
self._touch(existing.id)
except Exception:
logger.info("existing session %s unhealthy, recreating", existing.id[:12])
healthy = False
if healthy:
return existing
# Close unhealthy session (profile stays on disk)
await self.close(existing.id)
if existing_id:
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},
args=["--no-sandbox", "--disable-dev-shm-usage"],
)
# Grant clipboard access for the page origin
try:
parsed = urlparse(url)
origin = f"{parsed.scheme}://{parsed.netloc}"
await context.grant_permissions(["clipboard-read", "clipboard-write"], origin=origin)
except Exception:
logger.debug("clipboard permission grant failed (non-fatal)")
page = context.pages[0] if context.pages else await context.new_page()
session = BrowserSession(
id=uuid4().hex,
@@ -87,6 +116,9 @@ class BrowserSessionService:
)
self._sessions[session.id] = session
self._profiles[profile_key] = session.id
self._touch(session.id)
# Evict again after adding the new session so cap is enforced immediately
await self._evict_idle_sessions()
try:
await page.goto(url, wait_until="domcontentloaded", timeout=45000)
await self._autofill_login(page, login_config)
@@ -96,8 +128,13 @@ class BrowserSessionService:
raise
return session
def _touch(self, session_id: str) -> None:
"""Mark a session as recently active (reset idle timer)."""
self._last_event_at[session_id] = asyncio.get_event_loop().time()
async def screenshot(self, session_id: str) -> bytes:
session = self._get(session_id)
self._touch(session_id)
async with session.lock:
self._ensure_open(session)
return await session.page.screenshot(type="jpeg", quality=65, full_page=False)
@@ -111,6 +148,7 @@ class BrowserSessionService:
include_state: bool = True,
) -> dict[str, Any] | None:
session = self._get(session_id)
self._last_event_at[session_id] = asyncio.get_event_loop().time()
async with session.lock:
self._ensure_open(session)
page = session.page
@@ -156,12 +194,51 @@ class BrowserSessionService:
async def selected_text(self, session_id: str) -> str:
session = self._get(session_id)
self._touch(session_id)
async with session.lock:
self._ensure_open(session)
value = await session.page.evaluate("() => window.getSelection()?.toString() || ''")
return str(value or "")
async def read_clipboard(self, session_id: str) -> tuple[Optional[str], Optional[str]]:
"""Read the remote browser's clipboard text.
Returns (text, error_reason).
text is None when the clipboard is empty or unreadable.
error_reason is None on success or "empty" — non-None indicates a genuine failure.
"""
session = self._get(session_id)
self._touch(session_id)
async with session.lock:
self._ensure_open(session)
try:
result = await session.page.evaluate("""
async () => {
try {
const text = await navigator.clipboard.readText();
return text || null;
} catch (e) {
if (e instanceof DOMException) {
if (e.name === 'NotAllowedError') return 'ERROR:denied';
if (e.name === 'NotFoundError') return null;
}
return 'ERROR:' + (e.message || String(e));
}
}
""")
if result is None:
return None, None # empty clipboard
if isinstance(result, str) and result.startswith("ERROR:"):
reason = result[6:]
logger.debug("clipboard read error for %s: %s", session_id[:12], reason)
return None, reason
return str(result), None
except Exception as exc:
logger.warning("clipboard read failed for %s: %s", session_id[:12], exc)
return None, "read_failed"
async def close(self, session_id: str) -> None:
self._last_event_at.pop(session_id, None)
session = self._discard_session(session_id)
if not session:
return
@@ -185,6 +262,14 @@ class BrowserSessionService:
pass
async def shutdown(self) -> None:
# Cancel the background eviction loop
if self._evict_task is not None and not self._evict_task.done():
self._evict_task.cancel()
try:
await self._evict_task
except asyncio.CancelledError:
pass
self._evict_task = None
sessions = list(self._sessions)
for session_id in sessions:
await self.close(session_id)
@@ -194,6 +279,7 @@ class BrowserSessionService:
async def state(self, session_id: str) -> dict[str, Any]:
session = self._get(session_id)
self._touch(session_id)
async with session.lock:
self._ensure_open(session)
return await self._session_state(session)
@@ -217,6 +303,9 @@ class BrowserSessionService:
self._playwright = await async_playwright().start()
except Exception as exc:
raise BrowserDependencyError(f"Unable to start Playwright: {exc}") from exc
# Start background eviction loop
if self._evict_task is None or self._evict_task.done():
self._evict_task = asyncio.create_task(self._evict_loop())
async def _reset_page_zoom(self, session: BrowserSession) -> None:
try:
@@ -228,20 +317,38 @@ class BrowserSessionService:
except Exception:
pass
async def autofill_login(
self,
session_id: str,
login_config: Optional[dict[str, Any]],
) -> bool:
"""Public: manually trigger login autofill for an active session.
Only fills username/password fields — never auto-submits.
Returns True if fields were found and filled, False otherwise.
Never returns password data to the caller.
"""
session = self._get(session_id)
self._touch(session_id)
async with session.lock:
self._ensure_open(session)
return await self._autofill_login(session.page, login_config, max_wait_seconds=3.0, skip_submit=True)
async def _autofill_login(
self,
page: Any,
config: Optional[dict[str, Any]],
*,
max_wait_seconds: float = 8.0,
max_wait_seconds: float = 2.0,
poll_interval_seconds: float = 0.25,
) -> None:
skip_submit: bool = False,
) -> bool:
if not config or not config.get("enabled"):
return
return False
username = str(config.get("username") or "")
password = str(config.get("password") or "")
if not username or not password:
return
return False
try:
username_selectors = [
config.get("username_selector"),
@@ -268,17 +375,20 @@ class BrowserSessionService:
poll_interval_seconds=poll_interval_seconds,
)
if not username_locator or not password_locator:
logger.info("Login autofill skipped for %s: login fields not found", page.url)
return
logger.info("Login autofill skipped: login fields not found")
return False
await username_locator.fill(username, timeout=3000)
await password_locator.fill(password, timeout=3000)
submit_selector = str(config.get("submit_selector") or "").strip()
if submit_selector:
submit = await self._first_visible_locator(page, [submit_selector], timeout=500)
if submit:
await submit.click(timeout=3000)
if not skip_submit:
submit_selector = str(config.get("submit_selector") or "").strip()
if submit_selector:
submit = await self._first_visible_locator(page, [submit_selector], timeout=500)
if submit:
await submit.click(timeout=3000)
return True
except Exception as exc:
logger.info("Login autofill skipped for %s: %s", page.url, exc)
logger.info("Login autofill skipped: %s", exc)
return False
async def _wait_for_login_locators(
self,
@@ -345,6 +455,68 @@ class BrowserSessionService:
self._profiles.pop(session.profile_key, None)
return session
async def _evict_loop(self) -> None:
"""Background loop that runs every 5 minutes to evict idle sessions."""
while True:
await asyncio.sleep(300) # 5 minutes
try:
await self._evict_idle_sessions()
except Exception:
logger.exception("idle eviction loop error")
async def _evict_idle_sessions(self) -> None:
"""Close oldest idle sessions when over cap, or any past TTL."""
now = asyncio.get_event_loop().time()
# First: drop sessions past idle TTL (excluding just-created ones)
to_remove: list[str] = []
for sid, session in self._sessions.items():
if session.profile_key and session.profile_key.startswith("auth-capture-"):
continue # ephemeral sessions are handled separately
last_active = self._last_event_at.get(sid, 0.0)
if last_active > 0 and (now - last_active) > self.IDLE_TTL_SECONDS:
to_remove.append(sid)
for sid in to_remove:
logger.info("evicting idle session %s (no activity for >%ds)", sid[:12], self.IDLE_TTL_SECONDS)
await self.close(sid)
# Second: if still over cap, evict oldest by last_event_at
persistent = [(sid, s) for sid, s in self._sessions.items()
if not (s.profile_key or "").startswith("auth-capture-")]
if len(persistent) > self.MAX_SESSIONS:
persistent.sort(key=lambda x: self._last_event_at.get(x[0], 0.0))
excess = len(persistent) - self.MAX_SESSIONS
for sid, _ in persistent[:excess]:
logger.info("evicting session %s (over cap of %d)", sid[:12], self.MAX_SESSIONS)
await self.close(sid)
async def clear_profile(self, custom_page_id: int, url: str) -> None:
"""Close session for the page if active, then delete profile directory.
Raises RuntimeError if the directory cannot be fully removed.
"""
import shutil
# Close active session and use its profile_key (precise match)
profile_key: Optional[str] = None
try:
session = self.find_by_page_id(custom_page_id)
profile_key = session.profile_key
await self.close(session.id)
except KeyError:
pass
# Fallback: compute from URL (may be wrong if URL changed since session was created)
if not profile_key:
profile_key = self._profile_key(custom_page_id, url)
profile_dir = self._profile_dir(profile_key)
if profile_dir.exists():
shutil.rmtree(profile_dir) # no ignore_errors — let failure surface
if profile_dir.exists():
raise RuntimeError(
f"Failed to fully remove browser profile directory: {profile_dir}"
)
logger.info("cleared browser profile for page %d: %s", custom_page_id, profile_dir)
def _profile_dir(self, profile_key: str) -> Path:
root = Path(get_settings().browser_profiles_dir)
root.mkdir(parents=True, exist_ok=True)
@@ -383,6 +555,13 @@ class BrowserSessionService:
viewport={"width": width, "height": height},
args=["--no-sandbox", "--disable-dev-shm-usage"],
)
# Grant clipboard access for the page origin
try:
parsed = urlparse(url)
origin = f"{parsed.scheme}://{parsed.netloc}"
await context.grant_permissions(["clipboard-read", "clipboard-write"], origin=origin)
except Exception:
logger.debug("clipboard permission grant failed (non-fatal)")
page = context.pages[0] if context.pages else await context.new_page()
session = BrowserSession(
id=session_id,
@@ -394,6 +573,7 @@ class BrowserSessionService:
captured_headers=[],
)
self._sessions[session.id] = session
self._touch(session.id)
# Start CDP network capture BEFORE the initial page load,
# so we capture login redirects and auth headers from the start.
await self._start_cdp_capture(session)
+13
View File
@@ -62,6 +62,18 @@ def _check_upstream(upstream_id: int) -> None:
snapshot = build_snapshot(
upstream.id, upstream.base_url, upstream.api_prefix, groups, raw_rates
)
# ── Balance fetch (inside with block, client still open) ──
balance: Optional[float] = None
if upstream.balance_endpoint and upstream.balance_response_path:
try:
raw_balance = client.get_balance(upstream.balance_endpoint, upstream.balance_response_path)
if raw_balance is not None:
divisor = upstream.balance_divisor or 1.0
balance = raw_balance / divisor
except Exception as exc:
logger.warning("upstream %s balance fetch failed: %s", upstream.name, exc)
upstream.balance = balance
upstream.balance_updated_at = datetime.now(timezone.utc) if balance is not None else None
except Exception as exc:
# failure path
upstream.consecutive_failures = (upstream.consecutive_failures or 0) + 1
@@ -83,6 +95,7 @@ def _check_upstream(upstream_id: int) -> None:
return
# success path (client auto-closed by `with`)
prev_snapshot_row = (
db.query(UpstreamRateSnapshot)
.filter(UpstreamRateSnapshot.upstream_id == upstream_id)
+37 -2
View File
@@ -34,10 +34,18 @@ def _clean_auth_header_value(value: Any, field_name: str) -> str:
return ""
if text.startswith("Bearer "):
text = text[7:].strip()
# Try to sanitize non-latin-1 characters instead of hard-failing
try:
text.encode("latin-1")
except UnicodeEncodeError as exc:
raise UpstreamError(f"{field_name} contains non-HTTP-header characters; please re-extract and apply the full credential") from exc
except UnicodeEncodeError:
# Try stripping non-ASCII characters
cleaned = text.encode("ascii", errors="ignore").decode("ascii").strip()
if cleaned:
return cleaned
raise UpstreamError(
f"{field_name} 含有非 HTTP 标头字符(如中文或 emoji),"
f"请重新登录后再试"
) from None
return text
@@ -325,3 +333,30 @@ class UpstreamClient:
def get_group_rates(self, endpoint: str) -> Any:
return self._request("GET", endpoint)
def get_balance(self, endpoint: str, response_path: str) -> Optional[float]:
"""Call the balance endpoint and extract a numeric value using a dot-separated JSON path.
response_path 示例:
"balance" → resp["balance"]
"data.quota" → resp["data"]["quota"]
"data.total_balance" → resp["data"]["total_balance"]
"""
if not endpoint or not response_path:
return None
resp = self._request("GET", endpoint)
if not isinstance(resp, dict):
return None
parts = response_path.split(".")
value: Any = resp
for part in parts:
if isinstance(value, dict):
value = value.get(part)
else:
return None
if value is None:
return None
try:
return float(value)
except (ValueError, TypeError):
return None