Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 24 additions & 7 deletions src/api/admin.py
Original file line number Diff line number Diff line change
Expand Up @@ -716,16 +716,21 @@ class AddTokenRequest(BaseModel):
video_concurrency: int = -1
protocol_mode: Literal["protocol"] = "protocol"
google_cookies: str = Field(min_length=1, max_length=524288)
login_account: str = Field(min_length=3, max_length=254)
# 可选:留空时由 Cookie/会话自动识别,识别失败则使用占位邮箱
login_account: Optional[str] = Field(default=None, max_length=254)
login_password: Optional[str] = None
proxy_url: Optional[str] = None
auto_refresh_enabled: bool = True
refresh_interval_minutes: int = Field(default=30, ge=1, le=10080)

@field_validator("login_account")
@classmethod
def validate_login_account(cls, value: str) -> str:
def validate_login_account(cls, value: Optional[str]) -> Optional[str]:
if value is None:
return None
normalized = value.strip().lower()
if not normalized:
return None
if not re.fullmatch(r"[^\s@]+@[^\s@]+\.[^\s@]+", normalized):
raise ValueError("login_account must be a valid email address")
return normalized
Expand Down Expand Up @@ -880,13 +885,18 @@ class PluginTokenUpdateRequest(BaseModel):
model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)

google_cookies: str = Field(min_length=1, max_length=524288)
login_account: str = Field(min_length=3, max_length=254)
# 可选:留空时按 Cookie 会话定位账号,定位不到则使用占位邮箱
login_account: Optional[str] = Field(default=None, max_length=254)
protocol_mode: Literal["protocol"] = "protocol"

@field_validator("login_account")
@classmethod
def validate_login_account(cls, value: str) -> str:
def validate_login_account(cls, value: Optional[str]) -> Optional[str]:
if value is None:
return None
normalized = value.strip().lower()
if not normalized:
return None
if not re.fullmatch(r"[^\s@]+@[^\s@]+\.[^\s@]+", normalized):
raise ValueError("login_account must be a valid email address")
return normalized
Expand Down Expand Up @@ -2724,7 +2734,7 @@ async def plugin_update_token(
plugin_config = await db.get_plugin_config()

google_cookies = request.google_cookies
email = request.login_account
email = (request.login_account or "").strip().lower()

try:
result = await token_manager.flow_client.st_to_at(
Expand All @@ -2749,8 +2759,15 @@ async def plugin_update_token(
cookie_hash = hashlib.sha256(google_cookies.encode("utf-8")).hexdigest()[:32]
session_token = f"frontend-{cookie_hash}"

# Step 2: Check if token with this email exists
existing_token = await db.get_token_by_email(email)
# Step 2: 定位已有账号。邮箱可选:先用邮箱,未提供则按 Cookie 会话(st)反查,
# 都没有则生成与 add_token 一致的占位邮箱,避免“必须填 Google 账号”的硬性要求。
existing_token = await db.get_token_by_email(email) if email else None
if existing_token is None:
existing_token = await db.get_token_by_st(session_token)
if existing_token is not None and not email:
email = str(existing_token.email or existing_token.login_account or "")
if not email:
email = f"account-{cookie_hash[:16]}@flow.local"

if existing_token:
# Update existing token
Expand Down
22 changes: 20 additions & 2 deletions src/services/flow_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,21 @@
httpx = None


def _safe_updated_flow_cookies(owner, original) -> str:
"""读取 Cookie 最新快照,缺失实现时安全降级。

历史报错:'FlowClient' object has no attribute 'get_updated_flow_cookies'
成因是自定义 flow_frontend 覆盖了上游新增的 Cookie 作用域方法。
"""
getter = getattr(owner, "get_updated_flow_cookies", None)
if getter is None:
return str(original or "").strip()
try:
return getter(original)
except Exception:
return str(original or "").strip()


class _FlowClientBase(FlowFrontendMixin):
"""VideoFX API客户端"""

Expand Down Expand Up @@ -872,7 +887,10 @@ async def _make_text_request(
response = await session.post(url, **request_kwargs)

if cookie_jar is not None:
self._update_flow_cookie_jar(cookie_jar, response)
# 兼容自定义/旧版 flow_frontend:缺少该方法时跳过快照回写
updater = getattr(self, "_update_flow_cookie_jar", None)
if updater is not None:
updater(cookie_jar, response)

duration_ms = (time.time() - start_time) * 1000
response_text = response.text
Expand Down Expand Up @@ -4625,7 +4643,7 @@ async def _build_team_captcha_cookie_map(
cookie_storage = await self._resolve_flow_frontend_cookie_storage(
token_id=token_id
)
cookie_storage = self.get_updated_flow_cookies(cookie_storage)
cookie_storage = _safe_updated_flow_cookies(self, cookie_storage)
except Exception as e:
debug_logger.log_warning(
f"[reCAPTCHA {FLOW2API_TEAM_CAPTCHA_METHOD}] 读取账号 Cookie 失败: {e}"
Expand Down
18 changes: 17 additions & 1 deletion src/services/flow_current.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,22 @@
from ..core.logger import debug_logger
from .browser_cookie_utils import validate_flow_cookie_storage

def _safe_updated_flow_cookies(owner, original) -> str:
"""读取 Cookie 最新快照,缺失实现时安全降级。

旧版/自定义 flow_frontend 覆盖可能没有 get_updated_flow_cookies,
直接调用会抛 AttributeError(历史报错:
'FlowClient' object has no attribute 'get_updated_flow_cookies')。
"""
getter = getattr(owner, "get_updated_flow_cookies", None)
if getter is None:
return str(original or "").strip()
try:
return getter(original)
except Exception:
return str(original or "").strip()


class CurrentFlowClientMixin:
FRONTEND_ACCESS_TOKEN = "flow-frontend-cookie"
CURRENT_IMAGE_MODELS = frozenset({"BELUGA", "HARBOR_SEAL", "GEM_PIX_2"})
Expand Down Expand Up @@ -141,7 +157,7 @@ async def st_to_at(
"user": {},
"credits": credits.get("credits", 0),
"userPaygateTier": credits.get("userPaygateTier"),
"google_cookies": self.get_updated_flow_cookies(google_cookies),
"google_cookies": _safe_updated_flow_cookies(self, google_cookies),
}

async def create_project(
Expand Down
60 changes: 60 additions & 0 deletions tests/test_flow_cookie_snapshot_compat.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
"""Cookie 快照兼容性:缺失 get_updated_flow_cookies 时不能抛 AttributeError。

历史故障:自定义/旧版 flow_frontend 覆盖了上游新增的 Cookie 作用域方法,
导致添加 Google Cookie 时直接失败:

添加失败: Google Cookie 验证失败:
'FlowClient' object has no attribute 'get_updated_flow_cookies'

修复策略:调用点改为安全取值(getattr + 降级返回原始 Cookie 串),
既保留上游 Cookie 轮换快照能力,也能容忍不完整的 mixin 覆盖。
"""

import unittest

from src.services.flow_client import _safe_updated_flow_cookies as client_helper
from src.services.flow_current import _safe_updated_flow_cookies as current_helper

HELPERS = (client_helper, current_helper)


class _NoSuchMethod:
"""模拟缺少 get_updated_flow_cookies 的旧版 mixin。"""


class _WithMethod:
def get_updated_flow_cookies(self, original):
return f"{original};ROTATED=1"


class _RaisingMethod:
def get_updated_flow_cookies(self, original):
raise RuntimeError("boom")


class SafeUpdatedFlowCookiesTest(unittest.TestCase):
def test_missing_method_falls_back_to_original(self):
for helper in HELPERS:
with self.subTest(helper=helper.__module__):
self.assertEqual(helper(_NoSuchMethod(), "SID=abc"), "SID=abc")

def test_missing_method_normalizes_none(self):
for helper in HELPERS:
with self.subTest(helper=helper.__module__):
self.assertEqual(helper(_NoSuchMethod(), None), "")

def test_uses_method_when_present(self):
for helper in HELPERS:
with self.subTest(helper=helper.__module__):
self.assertEqual(
helper(_WithMethod(), "SID=abc"), "SID=abc;ROTATED=1"
)

def test_method_error_falls_back_to_original(self):
for helper in HELPERS:
with self.subTest(helper=helper.__module__):
self.assertEqual(helper(_RaisingMethod(), " SID=abc "), "SID=abc")


if __name__ == "__main__":
unittest.main()
Loading