From f38e9874789a6180b74ca920e7108f7ab6a43aea Mon Sep 17 00:00:00 2001 From: Wizard Date: Wed, 4 Nov 2026 00:41:02 -0400 Subject: [PATCH 1/3] =?UTF-8?q?feat(generation):=20=E9=85=8D=E9=A2=9D?= =?UTF-8?q?=E8=80=97=E5=B0=BD=E6=88=96=E9=A3=8E=E6=8E=A7=E6=8B=92=E7=BB=9D?= =?UTF-8?q?=E6=97=B6=E8=87=AA=E5=8A=A8=E5=88=87=E6=8D=A2=E8=B4=A6=E5=8F=B7?= =?UTF-8?q?=E9=87=8D=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 生成链路遇到 PUBLIC_ERROR_USER_QUOTA_REACHED / PUBLIC_ERROR_UNUSUAL_ACTIVITY 时排除当前 Token 并切换其他账号重试（最多 3 个账号）；提交结果不确定时不换号， 避免重复生成 - LoadBalancer 支持 exclude_token_ids，新增配额冷却（1h）与风控冷却（5min）， 冷却中的 Token 排到候选末尾 - 视频提交重试接入统一重试处理：每次被拒都通知打码服务自愈，UNUSUAL_ACTIVITY 同账号仅再重试一次并延迟 3s，其余交给上层换号 - 补充配额/风控换号与视频重试相关测试 Co-Authored-By: Claude Opus 5.5 --- src/services/flow_client.py | 3 + src/services/flow_current.py | 37 +- src/services/generation_handler.py | 667 ++++++++++++++++------------- src/services/load_balancer.py | 69 ++- tests/test_token_quota_failover.py | 154 +++++++ tests/test_veo_lite_support.py | 24 ++ 6 files changed, 650 insertions(+), 304 deletions(-) create mode 100644 tests/test_token_quota_failover.py diff --git a/src/services/flow_client.py b/src/services/flow_client.py index 8c634392..acb03e1f 100644 --- a/src/services/flow_client.py +++ b/src/services/flow_client.py @@ -4748,6 +4748,9 @@ def _get_retry_delay_seconds(self, error_str: str, retry_attempt: int) -> int: return 2 if "recaptcha" in error_lower: return 2 + if "public_error_unusual_activity" in error_lower: + # 风控拒绝后立刻重试几乎必然再次被拒，留出打码服务自愈时间 + return 3 return 1 def _resolve_recaptcha_runtime_settings( diff --git a/src/services/flow_current.py b/src/services/flow_current.py index 91a55c0a..3d73ad46 100644 --- a/src/services/flow_current.py +++ b/src/services/flow_current.py @@ -15,6 +15,10 @@ from ..core.logger import debug_logger from .browser_cookie_utils import validate_flow_cookie_storage +# 视频提交被 PUBLIC_ERROR_UNUSUAL_ACTIVITY 拒绝时，同账号最多再重试的次数； +# 之后交给上层切换账号（同账号同 IP 立刻重试大概率仍被风控）。 +UNUSUAL_ACTIVITY_SAME_TOKEN_RETRIES = 1 + def _safe_updated_flow_cookies(owner, original) -> str: """读取 Cookie 最新快照，缺失实现时安全降级。 @@ -548,6 +552,7 @@ async def _current_generate_video( model_keys.append(model_keys[0].replace("_relaxed", "")) attempt_budget = max(max_retries, len(model_keys)) fallback_model_key: Optional[str] = None + unusual_activity_hits = 0 last_error: Optional[Exception] = None for retry_attempt in range(attempt_budget): browser_id = None @@ -610,7 +615,37 @@ async def _current_generate_video( f"{request_model_key} -> {model_keys[-1]}" ) continue - if retry_attempt >= max_retries - 1: + retry_handler = getattr( + self, "_handle_retryable_generation_error", None + ) + if retry_handler is None: + if retry_attempt >= max_retries - 1: + raise + continue + if "public_error_unusual_activity" in error_text: + unusual_activity_hits += 1 + if ( + retry_attempt >= max_retries - 1 + or unusual_activity_hits > UNUSUAL_ACTIVITY_SAME_TOKEN_RETRIES + ): + # 不再同号重试，但仍需通知打码服务自愈（清缓存/重建浏览器） + await self._notify_browser_captcha_error( + browser_id=getattr(error, "browser_ref", None) or browser_id, + project_id=project_id, + error_reason=self._get_retry_reason(str(error)) + or str(error)[:120], + error_message=str(error), + ) + raise + should_retry = await retry_handler( + error, + retry_attempt, + max_retries, + browser_id, + project_id, + "[VIDEO] ", + ) + if not should_retry: raise finally: await self._notify_browser_captcha_request_finished(browser_id) diff --git a/src/services/generation_handler.py b/src/services/generation_handler.py index 67e53940..ba1cec3d 100644 --- a/src/services/generation_handler.py +++ b/src/services/generation_handler.py @@ -33,6 +33,10 @@ def _video_poll_attempt_budget( return attempts * (3 if upsample else 1) +# 账号配额耗尽时，单个请求最多尝试的账号数（含首个账号） +MAX_TOKEN_FAILOVER_ATTEMPTS = 3 + + # Model configuration MODEL_CONFIG = { # 图片生成 - GEM_PIX_2 (Gemini 3.0 Pro) @@ -1651,191 +1655,365 @@ async def handle_generation( progress=0, ) - # 2. 选择Token - debug_logger.log_info(f"[GENERATION] 正在选择可用Token...") - token_select_started_at = time.time() - - if generation_type == "image": - token = await self.load_balancer.select_token( - for_image_generation=True, - model=model, - reserve=False, - enforce_concurrency_filter=False, - track_pending=True, - ) - else: - token = await self.load_balancer.select_token( - for_video_generation=True, - model=model, - reserve=False, - enforce_concurrency_filter=False, - track_pending=True, - ) - perf_trace["token_select_ms"] = int( - (time.time() - token_select_started_at) * 1000 - ) + # 2. 选择Token（账号配额耗尽时自动切换到其他账号重试） + failed_token_ids: set = set() + last_failover_error: Optional[str] = None + for token_attempt in range(MAX_TOKEN_FAILOVER_ATTEMPTS): + token = None + generation_result = self._create_generation_result() + pending_token_state = {"active": False} + # 选择本次尝试的 Token + debug_logger.log_info(f"[GENERATION] 正在选择可用Token...") + token_select_started_at = time.time() - if not token: - error_msg = None - if self.load_balancer and hasattr( - self.load_balancer, "get_unavailable_reason" - ): - error_msg = await self.load_balancer.get_unavailable_reason( - for_image_generation=(generation_type == "image"), - for_video_generation=(generation_type == "video"), + if generation_type == "image": + token = await self.load_balancer.select_token( + for_image_generation=True, model=model, + reserve=False, + enforce_concurrency_filter=False, + track_pending=True, + exclude_token_ids=failed_token_ids, ) - if not error_msg: - error_msg = self._get_no_token_error_message(generation_type) - debug_logger.log_error(f"[GENERATION] {error_msg}") - record_generation_result( - generation_type, "no_token", time.time() - start_time - ) - await self._log_request( - token_id=None, - operation=request_operation, - request_data=request_payload, - response_data={"error": error_msg, "performance": perf_trace}, - status_code=503, - duration=time.time() - start_time, - log_id=request_log_state.get("id"), - status_text="failed", - progress=request_log_state.get("progress", 0), + else: + token = await self.load_balancer.select_token( + for_video_generation=True, + model=model, + reserve=False, + enforce_concurrency_filter=False, + track_pending=True, + exclude_token_ids=failed_token_ids, + ) + perf_trace["token_select_ms"] = int( + (time.time() - token_select_started_at) * 1000 ) - if stream: - yield self._create_stream_chunk(f"错误: {error_msg}\n") - yield self._create_error_response(error_msg, status_code=503) - return - - debug_logger.log_info(f"[GENERATION] 已选择Token: {token.id} ({token.email})") - pending_token_state["active"] = True - await self._update_request_log_progress( - request_log_state, - token_id=token.id, - status_text="token_selected", - progress=8, - response_extra={"token_email": token.email}, - ) - try: - # 3. 确保AT有效 - debug_logger.log_info(f"[GENERATION] 检查Token AT有效性...") - if stream: - yield self._create_stream_chunk("初始化生成环境...\n") - - await self._update_request_log_progress( - request_log_state, - token_id=token.id, - status_text="token_ready", - progress=15, - ) - ensure_at_started_at = time.time() - token = await self.token_manager.ensure_valid_token(token) - perf_trace["ensure_at_ms"] = int( - (time.time() - ensure_at_started_at) * 1000 - ) if not token: - error_msg = "Token AT无效或刷新失败" + error_msg = None + if self.load_balancer and hasattr( + self.load_balancer, "get_unavailable_reason" + ): + error_msg = await self.load_balancer.get_unavailable_reason( + for_image_generation=(generation_type == "image"), + for_video_generation=(generation_type == "video"), + model=model, + ) + if not error_msg: + error_msg = self._get_no_token_error_message(generation_type) + if last_failover_error: + error_msg = ( + f"{last_failover_error}（已尝试 {len(failed_token_ids)} 个账号，" + f"没有其他可切换的账号）" + ) debug_logger.log_error(f"[GENERATION] {error_msg}") record_generation_result( - generation_type, "failed", time.time() - start_time + generation_type, "no_token", time.time() - start_time ) - if stream: - yield self._create_stream_chunk(f"错误: {error_msg}\n") - yield self._create_error_response(error_msg, status_code=503) - return - - # 4. 确保Project存在 - debug_logger.log_info(f"[GENERATION] 检查/创建Project...") - - if not supports_model_for_tier(model, token.user_paygate_tier): - required_tier = get_required_paygate_tier_for_model(model) - error_msg = ( - "当前模型需要 " - + get_paygate_tier_label(required_tier) - + " 账号: " - + model - ) - debug_logger.log_error(f"[GENERATION] {error_msg}") - record_generation_result( - generation_type, "failed", time.time() - start_time + await self._log_request( + token_id=None, + operation=request_operation, + request_data=request_payload, + response_data={"error": error_msg, "performance": perf_trace}, + status_code=503, + duration=time.time() - start_time, + log_id=request_log_state.get("id"), + status_text="failed", + progress=request_log_state.get("progress", 0), ) if stream: yield self._create_stream_chunk(f"错误: {error_msg}\n") - yield self._create_error_response(error_msg, status_code=403) + yield self._create_error_response(error_msg, status_code=503) return - ensure_project_started_at = time.time() - project_id = await self.token_manager.ensure_project_exists(token.id) - perf_trace["ensure_project_ms"] = int( - (time.time() - ensure_project_started_at) * 1000 - ) - debug_logger.log_info(f"[GENERATION] Project ID: {project_id}") + debug_logger.log_info(f"[GENERATION] 已选择Token: {token.id} ({token.email})") + pending_token_state["active"] = True await self._update_request_log_progress( request_log_state, token_id=token.id, - status_text="project_ready", - progress=22, - response_extra={"project_id": project_id}, - ) - prefill_action = ( - "IMAGE_GENERATION" if generation_type == "image" else "VIDEO_GENERATION" - ) - await self.flow_client.prefill_remote_browser_pool( - project_id=project_id, - action=prefill_action, - token_id=token.id, + status_text="token_selected", + progress=8, + response_extra={"token_email": token.email}, ) - # 5. 根据类型处理 - generation_pipeline_started_at = time.time() - if generation_type == "image": - debug_logger.log_info(f"[GENERATION] 开始图片生成流程...") - async for chunk in self._handle_image_generation( - token, - project_id, - model_config, - prompt, - images, - stream, - perf_trace=perf_trace, - generation_result=generation_result, - response_state=response_state, - request_log_state=request_log_state, - pending_token_state=pending_token_state, - ): - yield chunk - else: # video - debug_logger.log_info(f"[GENERATION] 开始视频生成流程...") - async for chunk in self._handle_video_generation( - token, - project_id, - model_config, - prompt, - images, - stream, - perf_trace=perf_trace, - generation_result=generation_result, - response_state=response_state, - request_log_state=request_log_state, - pending_token_state=pending_token_state, - video_media_id=video_media_id, - ): - yield chunk - perf_trace["generation_pipeline_ms"] = int( - (time.time() - generation_pipeline_started_at) * 1000 - ) + try: + # 3. 确保AT有效 + debug_logger.log_info(f"[GENERATION] 检查Token AT有效性...") + if stream: + yield self._create_stream_chunk("初始化生成环境...\n") + + await self._update_request_log_progress( + request_log_state, + token_id=token.id, + status_text="token_ready", + progress=15, + ) + ensure_at_started_at = time.time() + token = await self.token_manager.ensure_valid_token(token) + perf_trace["ensure_at_ms"] = int( + (time.time() - ensure_at_started_at) * 1000 + ) + if not token: + error_msg = "Token AT无效或刷新失败" + debug_logger.log_error(f"[GENERATION] {error_msg}") + record_generation_result( + generation_type, "failed", time.time() - start_time + ) + if stream: + yield self._create_stream_chunk(f"错误: {error_msg}\n") + yield self._create_error_response(error_msg, status_code=503) + return + + # 4. 确保Project存在 + debug_logger.log_info(f"[GENERATION] 检查/创建Project...") + + if not supports_model_for_tier(model, token.user_paygate_tier): + required_tier = get_required_paygate_tier_for_model(model) + error_msg = ( + "当前模型需要 " + + get_paygate_tier_label(required_tier) + + " 账号: " + + model + ) + debug_logger.log_error(f"[GENERATION] {error_msg}") + record_generation_result( + generation_type, "failed", time.time() - start_time + ) + if stream: + yield self._create_stream_chunk(f"错误: {error_msg}\n") + yield self._create_error_response(error_msg, status_code=403) + return + + ensure_project_started_at = time.time() + project_id = await self.token_manager.ensure_project_exists(token.id) + perf_trace["ensure_project_ms"] = int( + (time.time() - ensure_project_started_at) * 1000 + ) + debug_logger.log_info(f"[GENERATION] Project ID: {project_id}") + await self._update_request_log_progress( + request_log_state, + token_id=token.id, + status_text="project_ready", + progress=22, + response_extra={"project_id": project_id}, + ) + prefill_action = ( + "IMAGE_GENERATION" if generation_type == "image" else "VIDEO_GENERATION" + ) + await self.flow_client.prefill_remote_browser_pool( + project_id=project_id, + action=prefill_action, + token_id=token.id, + ) + + # 5. 根据类型处理 + generation_pipeline_started_at = time.time() + if generation_type == "image": + debug_logger.log_info(f"[GENERATION] 开始图片生成流程...") + async for chunk in self._handle_image_generation( + token, + project_id, + model_config, + prompt, + images, + stream, + perf_trace=perf_trace, + generation_result=generation_result, + response_state=response_state, + request_log_state=request_log_state, + pending_token_state=pending_token_state, + ): + yield chunk + else: # video + debug_logger.log_info(f"[GENERATION] 开始视频生成流程...") + async for chunk in self._handle_video_generation( + token, + project_id, + model_config, + prompt, + images, + stream, + perf_trace=perf_trace, + generation_result=generation_result, + response_state=response_state, + request_log_state=request_log_state, + pending_token_state=pending_token_state, + video_media_id=video_media_id, + ): + yield chunk + perf_trace["generation_pipeline_ms"] = int( + (time.time() - generation_pipeline_started_at) * 1000 + ) + + # 6. 记录使用 + if not generation_result.get("success"): + error_msg = generation_result.get("error_message") or "生成未成功完成" + debug_logger.log_warning( + f"[GENERATION] 生成未成功，不扣次数: {error_msg}" + ) + if token: + await self.token_manager.record_error(token.id) + duration = time.time() - start_time + record_generation_result(generation_type, "failed", duration) + perf_trace["status"] = "failed" + perf_trace["total_ms"] = int(duration * 1000) + perf_trace["error"] = error_msg + prompt_for_log = ( + prompt if len(prompt) <= 2000 else f"{prompt[:2000]}...(truncated)" + ) + await self._log_request( + token.id if token else None, + request_operation, + request_payload, + {"error": error_msg, "performance": perf_trace}, + 500, + duration, + log_id=request_log_state.get("id"), + status_text="failed", + progress=request_log_state.get("progress", 0), + ) + if not generation_result.get("error_emitted"): + if stream: + yield self._create_stream_chunk(f"错误: {error_msg}\n") + yield self._create_error_response(error_msg, status_code=500) + return + + is_video = generation_type == "video" + await self.token_manager.record_usage(token.id, is_video=is_video) + + # 重置错误计数 (请求成功时清空连续错误计数) + await self.token_manager.record_success(token.id) + self.load_balancer.clear_quota_cooldown( + token.id, "video" if is_video else "image" + ) + + debug_logger.log_info(f"[GENERATION] ✅ 生成成功完成") + + # 7. 记录成功日志 + duration = time.time() - start_time + record_generation_result(generation_type, "success", duration) + perf_trace["status"] = "success" + perf_trace["total_ms"] = int(duration * 1000) + # 日志中保留更完整的 prompt，避免管理页只看到过短内容 + prompt_for_log = ( + prompt if len(prompt) <= 2000 else f"{prompt[:2000]}...(truncated)" + ) + + # 构建响应数据，包含生成的URL + response_data = { + "status": "success", + "model": model, + "prompt": prompt_for_log, + "performance": perf_trace, + } + + # 添加生成的URL（如果有） + if response_state.get("url"): + response_data["url"] = response_state["url"] + if response_state.get("generated_assets"): + response_data["generated_assets"] = response_state["generated_assets"] + image_perf = ( + perf_trace.get("image_generation", {}) + if isinstance(perf_trace, dict) + else {} + ) + video_perf = ( + perf_trace.get("video_generation", {}) + if isinstance(perf_trace, dict) + else {} + ) + debug_logger.log_info( + f"[PERF] [{request_id}] total={perf_trace.get('total_ms', 0)}ms, " + f"select={perf_trace.get('token_select_ms', 0)}ms, " + f"ensure_at={perf_trace.get('ensure_at_ms', 0)}ms, " + f"project={perf_trace.get('ensure_project_ms', 0)}ms, " + f"pipeline={perf_trace.get('generation_pipeline_ms', 0)}ms, " + f"slot_wait={image_perf.get('slot_wait_ms', 0)}ms, " + f"launch_queue={image_perf.get('launch_queue_wait_ms', 0)}ms, " + f"launch_stagger={image_perf.get('launch_stagger_wait_ms', 0)}ms, " + f"video_slot_wait={video_perf.get('slot_wait_ms', 0)}ms" + ) - # 6. 记录使用 - if not generation_result.get("success"): - error_msg = generation_result.get("error_message") or "生成未成功完成" - debug_logger.log_warning( - f"[GENERATION] 生成未成功，不扣次数: {error_msg}" + await self._log_request( + token.id, + request_operation, + request_payload, + response_data, + 200, + duration, + log_id=request_log_state.get("id"), + status_text="completed", + progress=100, + ) + + except asyncio.CancelledError: + error_msg = "生成已取消: 客户端连接已断开" + debug_logger.log_warning(f"[GENERATION] ⚠️ {error_msg}") + duration = time.time() - start_time + record_generation_result( + generation_type or "unknown", "cancelled", duration + ) + perf_trace["status"] = "failed" + perf_trace["total_ms"] = int(duration * 1000) + perf_trace["error"] = error_msg + prompt_for_log = ( + prompt if len(prompt) <= 2000 else f"{prompt[:2000]}...(truncated)" + ) + await self._log_request( + token.id if token else None, + request_operation if generation_type else "generate_unknown", + request_payload if "request_payload" in locals() else {"model": model}, + {"error": error_msg, "performance": perf_trace}, + 499, + duration, + log_id=request_log_state.get("id"), + status_text="failed", + progress=request_log_state.get("progress", 0), ) + raise + except Exception as e: + error_msg = f"生成失败: {str(e)}" + debug_logger.log_error(f"[GENERATION] 生成失败: {error_msg}") + is_quota_error = self._is_quota_exhausted_error(e) + if ( + token + and (is_quota_error or self._is_unusual_activity_error(e)) + # 提交结果不确定时换号可能导致重复生成 + and not getattr(e, "submission_uncertain", False) + and not getattr(e, "outcome_unknown", False) + ): + scope = "image" if generation_type == "image" else "video" + if is_quota_error: + self.load_balancer.mark_quota_cooldown(token.id, scope) + failover_reason = "配额耗尽" + failover_notice = "当前账号配额已用尽，正在切换其他账号重试...\n" + else: + self.load_balancer.mark_risk_cooldown(token.id, scope) + failover_reason = "触发风控" + failover_notice = "当前账号触发风控，正在切换其他账号重试...\n" + failed_token_ids.add(token.id) + last_failover_error = error_msg + if token_attempt < MAX_TOKEN_FAILOVER_ATTEMPTS - 1: + perf_trace.setdefault("token_failovers", []).append( + {"token_id": token.id, "error": str(e)[:200]} + ) + debug_logger.log_warning( + f"[GENERATION] Token {token.id} {failover_reason}，切换账号重试 " + f"({token_attempt + 2}/{MAX_TOKEN_FAILOVER_ATTEMPTS})" + ) + if stream: + yield self._create_stream_chunk(failover_notice) + continue if token: - await self.token_manager.record_error(token.id) + if self._should_count_token_error(e): + await self.token_manager.record_error(token.id) + else: + debug_logger.log_info( + f"[GENERATION] 跳过 token 错误计数: token_id={token.id}, reason={str(e)[:200]}" + ) + + # 先将最终失败状态落库，再返回错误响应，避免日志停在 102。 duration = time.time() - start_time - record_generation_result(generation_type, "failed", duration) + record_generation_result(generation_type or "unknown", "failed", duration) perf_trace["status"] = "failed" perf_trace["total_ms"] = int(duration * 1000) perf_trace["error"] = error_msg @@ -1844,8 +2022,8 @@ async def handle_generation( ) await self._log_request( token.id if token else None, - request_operation, - request_payload, + request_operation if generation_type else "generate_unknown", + request_payload if "request_payload" in locals() else {"model": model}, {"error": error_msg, "performance": perf_trace}, 500, duration, @@ -1853,144 +2031,20 @@ async def handle_generation( status_text="failed", progress=request_log_state.get("progress", 0), ) - if not generation_result.get("error_emitted"): - if stream: - yield self._create_stream_chunk(f"错误: {error_msg}\n") - yield self._create_error_response(error_msg, status_code=500) - return - - is_video = generation_type == "video" - await self.token_manager.record_usage(token.id, is_video=is_video) - - # 重置错误计数 (请求成功时清空连续错误计数) - await self.token_manager.record_success(token.id) - - debug_logger.log_info(f"[GENERATION] ✅ 生成成功完成") - - # 7. 记录成功日志 - duration = time.time() - start_time - record_generation_result(generation_type, "success", duration) - perf_trace["status"] = "success" - perf_trace["total_ms"] = int(duration * 1000) - # 日志中保留更完整的 prompt，避免管理页只看到过短内容 - prompt_for_log = ( - prompt if len(prompt) <= 2000 else f"{prompt[:2000]}...(truncated)" - ) - - # 构建响应数据，包含生成的URL - response_data = { - "status": "success", - "model": model, - "prompt": prompt_for_log, - "performance": perf_trace, - } - - # 添加生成的URL（如果有） - if response_state.get("url"): - response_data["url"] = response_state["url"] - if response_state.get("generated_assets"): - response_data["generated_assets"] = response_state["generated_assets"] - image_perf = ( - perf_trace.get("image_generation", {}) - if isinstance(perf_trace, dict) - else {} - ) - video_perf = ( - perf_trace.get("video_generation", {}) - if isinstance(perf_trace, dict) - else {} - ) - debug_logger.log_info( - f"[PERF] [{request_id}] total={perf_trace.get('total_ms', 0)}ms, " - f"select={perf_trace.get('token_select_ms', 0)}ms, " - f"ensure_at={perf_trace.get('ensure_at_ms', 0)}ms, " - f"project={perf_trace.get('ensure_project_ms', 0)}ms, " - f"pipeline={perf_trace.get('generation_pipeline_ms', 0)}ms, " - f"slot_wait={image_perf.get('slot_wait_ms', 0)}ms, " - f"launch_queue={image_perf.get('launch_queue_wait_ms', 0)}ms, " - f"launch_stagger={image_perf.get('launch_stagger_wait_ms', 0)}ms, " - f"video_slot_wait={video_perf.get('slot_wait_ms', 0)}ms" - ) - - await self._log_request( - token.id, - request_operation, - request_payload, - response_data, - 200, - duration, - log_id=request_log_state.get("id"), - status_text="completed", - progress=100, - ) - - except asyncio.CancelledError: - error_msg = "生成已取消: 客户端连接已断开" - debug_logger.log_warning(f"[GENERATION] ⚠️ {error_msg}") - duration = time.time() - start_time - record_generation_result( - generation_type or "unknown", "cancelled", duration - ) - perf_trace["status"] = "failed" - perf_trace["total_ms"] = int(duration * 1000) - perf_trace["error"] = error_msg - prompt_for_log = ( - prompt if len(prompt) <= 2000 else f"{prompt[:2000]}...(truncated)" - ) - await self._log_request( - token.id if token else None, - request_operation if generation_type else "generate_unknown", - request_payload if "request_payload" in locals() else {"model": model}, - {"error": error_msg, "performance": perf_trace}, - 499, - duration, - log_id=request_log_state.get("id"), - status_text="failed", - progress=request_log_state.get("progress", 0), - ) - raise - except Exception as e: - error_msg = f"生成失败: {str(e)}" - debug_logger.log_error(f"[GENERATION] 生成失败: {error_msg}") - if token: - if self._should_count_token_error(e): - await self.token_manager.record_error(token.id) - else: - debug_logger.log_info( - f"[GENERATION] 跳过 token 错误计数: token_id={token.id}, reason={str(e)[:200]}" + if stream: + yield self._create_stream_chunk(f"错误: {error_msg}\n") + yield self._create_error_response(error_msg, status_code=500) + finally: + if pending_token_state.get("active") and token and self.load_balancer: + await self.load_balancer.release_pending( + token.id, + for_image_generation=(generation_type == "image"), + for_video_generation=(generation_type == "video"), ) + pending_token_state["active"] = False - # 先将最终失败状态落库，再返回错误响应，避免日志停在 102。 - duration = time.time() - start_time - record_generation_result(generation_type or "unknown", "failed", duration) - perf_trace["status"] = "failed" - perf_trace["total_ms"] = int(duration * 1000) - perf_trace["error"] = error_msg - prompt_for_log = ( - prompt if len(prompt) <= 2000 else f"{prompt[:2000]}...(truncated)" - ) - await self._log_request( - token.id if token else None, - request_operation if generation_type else "generate_unknown", - request_payload if "request_payload" in locals() else {"model": model}, - {"error": error_msg, "performance": perf_trace}, - 500, - duration, - log_id=request_log_state.get("id"), - status_text="failed", - progress=request_log_state.get("progress", 0), - ) - if stream: - yield self._create_stream_chunk(f"错误: {error_msg}\n") - yield self._create_error_response(error_msg, status_code=500) - finally: - if pending_token_state.get("active") and token and self.load_balancer: - await self.load_balancer.release_pending( - token.id, - for_image_generation=(generation_type == "image"), - for_video_generation=(generation_type == "video"), - ) - pending_token_state["active"] = False + # 本次尝试已成功或以非切换类错误结束 + break def _get_no_token_error_message(self, generation_type: str) -> str: """获取无可用Token时的详细错误信息""" @@ -1999,6 +2053,16 @@ def _get_no_token_error_message(self, generation_type: str) -> str: else: return "没有可用的Token进行视频生成。所有Token都处于禁用、冷却、配额耗尽或已过期状态。" + @staticmethod + def _is_quota_exhausted_error(error: Exception) -> bool: + """账号配额耗尽：换一个账号通常可以成功。""" + return "public_error_user_quota_reached" in str(error or "").lower() + + @staticmethod + def _is_unusual_activity_error(error: Exception) -> bool: + """账号/IP 被风控拒绝：提交被拒、未创建任务，换一个账号通常可以成功。""" + return "public_error_unusual_activity" in str(error or "").lower() + def _should_count_token_error(self, error: Exception) -> bool: """判断失败是否应计入 token 连续错误。 @@ -2015,6 +2079,7 @@ def _should_count_token_error(self, error: Exception) -> bool: "recaptcha 验证失败", "recaptcha 错误", "public_error_unusual_activity", + "public_error_user_quota_reached", "too much traffic", "error_no_slot_available", "打码服务资源不足", @@ -2799,7 +2864,7 @@ async def _handle_video_generation( await self.token_manager.record_usage(token.id, is_video=True) await self.token_manager.record_success(token.id) if getattr(self, "load_balancer", None): - self.load_balancer.clear_quota_cooldown(token.id, model_key) + self.load_balancer.clear_quota_cooldown(token.id, "video") if self.proxy_manager and hasattr( self.proxy_manager, "record_attempt_success", diff --git a/src/services/load_balancer.py b/src/services/load_balancer.py index 88627e69..1f1f0bc1 100644 --- a/src/services/load_balancer.py +++ b/src/services/load_balancer.py @@ -2,7 +2,8 @@ import asyncio import random -from typing import Optional, Dict +import time +from typing import Optional, Dict, Iterable from ..core.models import Token from ..core.config import config from ..core.account_tiers import ( @@ -14,6 +15,11 @@ from .concurrency_manager import ConcurrencyManager from ..core.logger import debug_logger +# 配额耗尽 (PUBLIC_ERROR_USER_QUOTA_REACHED) 后降低该 token 优先级的时长 +QUOTA_COOLDOWN_SECONDS = 3600 +# 风控拒绝 (PUBLIC_ERROR_UNUSUAL_ACTIVITY) 后降低该 token 优先级的时长 +RISK_COOLDOWN_SECONDS = 300 + class LoadBalancer: """Token load balancer with load-aware selection""" @@ -32,6 +38,53 @@ def __init__( "default": None, } self._rr_lock = asyncio.Lock() + # (token_id, "image"/"video") -> monotonic expiry; quota-exhausted tokens sort last + self._quota_cooldown: Dict[tuple[int, str], float] = {} + + @staticmethod + def _scope(for_image_generation: bool, for_video_generation: bool) -> Optional[str]: + if for_image_generation: + return "image" + if for_video_generation: + return "video" + return None + + def mark_quota_cooldown( + self, token_id: int, scope: str, seconds: float = QUOTA_COOLDOWN_SECONDS + ): + """配额耗尽后让该 token 在一段时间内排在候选末尾（不剔除，避免全部冷却时无号可用）。""" + self._quota_cooldown[(token_id, scope)] = time.monotonic() + seconds + debug_logger.log_warning( + f"[LOAD_BALANCER] Token {token_id} {scope} 配额耗尽，冷却 {int(seconds)}s 内优先使用其他账号" + ) + + def mark_risk_cooldown( + self, token_id: int, scope: str, seconds: float = RISK_COOLDOWN_SECONDS + ): + """风控拒绝后短暂让该 token 排在候选末尾；不会缩短已有的配额冷却。""" + key = (token_id, scope) + expires_at = time.monotonic() + seconds + if self._quota_cooldown.get(key, 0) < expires_at: + self._quota_cooldown[key] = expires_at + debug_logger.log_warning( + f"[LOAD_BALANCER] Token {token_id} {scope} 触发风控，冷却 {int(seconds)}s 内优先使用其他账号" + ) + + def clear_quota_cooldown(self, token_id: int, scope: Optional[str] = None): + for key in list(self._quota_cooldown): + if key[0] == token_id and (scope is None or key[1] == scope): + self._quota_cooldown.pop(key, None) + + def _in_quota_cooldown(self, token_id: int, scope: Optional[str]) -> bool: + if scope is None: + return False + expires_at = self._quota_cooldown.get((token_id, scope)) + if expires_at is None: + return False + if time.monotonic() >= expires_at: + self._quota_cooldown.pop((token_id, scope), None) + return False + return True async def _get_pending_count( self, token_id: int, for_image_generation: bool, for_video_generation: bool @@ -173,6 +226,7 @@ async def select_token( reserve: bool = False, enforce_concurrency_filter: bool = True, track_pending: bool = False, + exclude_token_ids: Optional[Iterable[int]] = None, ) -> Optional[Token]: """ Select a token using load-aware balancing @@ -189,6 +243,7 @@ async def select_token( track_pending: Whether to count the selected token as a queued request immediately. This smooths burst distribution before the hard concurrency slot is acquired. + exclude_token_ids: Tokens to skip (e.g. ones that already failed this request). Returns: Selected token or None if no available tokens @@ -210,8 +265,13 @@ async def select_token( available_tokens = [] filtered_reasons = {} required_tier = get_required_paygate_tier_for_model(model) + excluded = set(exclude_token_ids or ()) + scope = self._scope(for_image_generation, for_video_generation) for token in active_tokens: + if token.id in excluded: + filtered_reasons[token.id] = "本次请求已失败过" + continue normalized_tier = normalize_user_paygate_tier(token.user_paygate_tier) if model and not supports_model_for_tier(model, normalized_tier): filtered_reasons[token.id] = ( @@ -265,6 +325,7 @@ async def select_token( "inflight": inflight, "remaining": remaining, "needs_refresh": self.token_manager.needs_at_refresh(token), + "quota_cooldown": self._in_quota_cooldown(token.id, scope), "random": random.random(), } ) @@ -321,6 +382,9 @@ async def select_token( if ready_candidates and refresh_candidates: available_tokens = ready_candidates + refresh_candidates + # 配额冷却中的 token 排到最后（稳定排序，保持上面的顺序） + available_tokens.sort(key=lambda item: 1 if item["quota_cooldown"] else 0) + debug_logger.log_info("[LOAD_BALANCER] 候选Token负载:") for item in available_tokens: token = item["token"] @@ -328,7 +392,8 @@ async def select_token( debug_logger.log_info( f"[LOAD_BALANCER] - Token {token.id} ({token.email}) " f"inflight={item['inflight']}, remaining={remaining}, " - f"needs_refresh={item['needs_refresh']}, credits={token.credits}" + f"needs_refresh={item['needs_refresh']}, " + f"quota_cooldown={item['quota_cooldown']}, credits={token.credits}" ) # 只为候选列表中真正尝试到的 token 做 AT 校验，避免每次请求把所有 token 全扫一遍 diff --git a/tests/test_token_quota_failover.py b/tests/test_token_quota_failover.py new file mode 100644 index 00000000..6d6d2706 --- /dev/null +++ b/tests/test_token_quota_failover.py @@ -0,0 +1,154 @@ +import json +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +from src.services.generation_handler import ( + MAX_TOKEN_FAILOVER_ATTEMPTS, + GenerationHandler, +) +from src.services.load_balancer import LoadBalancer + + +MODEL = "veo-3.1-fast-i2v-6s-landscape" +QUOTA_ERROR = "Flow frontend RPC rejected: rpc=eb1hJf, error=PUBLIC_ERROR_USER_QUOTA_REACHED" +UNUSUAL_ERROR = "Flow frontend RPC rejected: rpc=eb1hJf, error=PUBLIC_ERROR_UNUSUAL_ACTIVITY" + + +def _token(token_id): + return SimpleNamespace( + id=token_id, + email=f"user{token_id}@example.com", + user_paygate_tier="PAYGATE_TIER_TWO", + image_enabled=True, + video_enabled=True, + credits=100, + image_concurrency=-1, + video_concurrency=-1, + ) + + +def _build_handler(tokens, quota_token_ids, error=QUOTA_ERROR): + token_manager = MagicMock() + token_manager.get_active_tokens = AsyncMock(return_value=tokens) + token_manager.ensure_valid_token = AsyncMock(side_effect=lambda t: t) + token_manager.ensure_project_exists = AsyncMock(return_value="project-1") + token_manager.needs_at_refresh = MagicMock(return_value=False) + token_manager.record_error = AsyncMock() + token_manager.record_usage = AsyncMock() + token_manager.record_success = AsyncMock() + + load_balancer = LoadBalancer(token_manager) + load_balancer._check_extension_route = AsyncMock(return_value=(True, "")) + + handler = GenerationHandler.__new__(GenerationHandler) + handler.token_manager = token_manager + handler.load_balancer = load_balancer + handler.flow_client = MagicMock() + handler.flow_client.prefill_remote_browser_pool = AsyncMock() + handler._log_request = AsyncMock(return_value=1) + handler._update_request_log_progress = AsyncMock() + + attempted = [] + + async def fake_video(token, *args, generation_result=None, **kwargs): + attempted.append(token.id) + if token.id in quota_token_ids: + raise error() if callable(error) else Exception(error) + handler._mark_generation_succeeded(generation_result) + yield handler._create_completion_response("http://video", media_type="video") + + handler._handle_video_generation = fake_video + return handler, attempted + + +async def _collect(handler): + with patch("src.services.generation_handler.record_generation_result"): + return [chunk async for chunk in handler.handle_generation(MODEL, "prompt")] + + +class TokenQuotaFailoverTests(unittest.IsolatedAsyncioTestCase): + async def test_switches_to_another_token_on_quota_error(self): + handler, attempted = _build_handler([_token(1), _token(2)], {1}) + + # 固定随机打散顺序，让 token 1 先被选中 + with patch("src.services.load_balancer.random.random", side_effect=[0.1, 0.9] * 4): + chunks = await _collect(handler) + + self.assertEqual(attempted, [1, 2]) + self.assertIn("http://video", chunks[-1]) + self.assertNotIn("error", json.loads(chunks[-1])) + # quota errors must not push the account toward auto-disable + handler.token_manager.record_error.assert_not_awaited() + self.assertTrue(handler.load_balancer._in_quota_cooldown(1, "video")) + self.assertFalse(handler.load_balancer._in_quota_cooldown(2, "video")) + + async def test_gives_up_after_max_attempts(self): + tokens = [_token(i) for i in range(1, 6)] + handler, attempted = _build_handler(tokens, {t.id for t in tokens}) + + chunks = await _collect(handler) + + self.assertEqual(len(attempted), MAX_TOKEN_FAILOVER_ATTEMPTS) + self.assertEqual(len(set(attempted)), MAX_TOKEN_FAILOVER_ATTEMPTS) + error = json.loads(chunks[-1])["error"] + self.assertIn("PUBLIC_ERROR_USER_QUOTA_REACHED", error["message"]) + self.assertEqual(handler._log_request.await_args.args[4], 500) + + async def test_reports_quota_error_when_no_token_left(self): + handler, attempted = _build_handler([_token(1), _token(2)], {1, 2}) + + chunks = await _collect(handler) + + self.assertEqual(sorted(attempted), [1, 2]) + error = json.loads(chunks[-1])["error"] + self.assertIn("PUBLIC_ERROR_USER_QUOTA_REACHED", error["message"]) + self.assertIn("已尝试 2 个账号", error["message"]) + + async def test_cooled_down_token_is_tried_last(self): + handler, attempted = _build_handler([_token(1), _token(2)], set()) + handler.load_balancer.mark_quota_cooldown(1, "video") + + with patch("src.services.load_balancer.random.random", side_effect=[0.1, 0.9] * 4): + await _collect(handler) + + self.assertEqual(attempted, [2]) + + +class TokenUnusualActivityFailoverTests(unittest.IsolatedAsyncioTestCase): + async def test_switches_to_another_token_on_unusual_activity(self): + handler, attempted = _build_handler([_token(1), _token(2)], {1}, UNUSUAL_ERROR) + + with patch("src.services.load_balancer.random.random", side_effect=[0.1, 0.9] * 4): + chunks = await _collect(handler) + + self.assertEqual(attempted, [1, 2]) + self.assertNotIn("error", json.loads(chunks[-1])) + handler.token_manager.record_error.assert_not_awaited() + # 风控冷却较短，且同样让该 token 排到后面 + self.assertTrue(handler.load_balancer._in_quota_cooldown(1, "video")) + + async def test_risk_cooldown_does_not_shorten_quota_cooldown(self): + handler, _ = _build_handler([_token(1)], set()) + lb = handler.load_balancer + lb.mark_quota_cooldown(1, "video") + quota_expiry = lb._quota_cooldown[(1, "video")] + lb.mark_risk_cooldown(1, "video") + self.assertEqual(lb._quota_cooldown[(1, "video")], quota_expiry) + + async def test_uncertain_submission_is_not_failed_over(self): + def uncertain(): + error = Exception(UNUSUAL_ERROR) + error.submission_uncertain = True + return error + + handler, attempted = _build_handler([_token(1), _token(2)], {1, 2}, uncertain) + + chunks = await _collect(handler) + + self.assertEqual(len(attempted), 1) + self.assertIn("PUBLIC_ERROR_UNUSUAL_ACTIVITY", json.loads(chunks[-1])["error"]["message"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_veo_lite_support.py b/tests/test_veo_lite_support.py index 8dca5839..99261295 100644 --- a/tests/test_veo_lite_support.py +++ b/tests/test_veo_lite_support.py @@ -361,6 +361,30 @@ async def test_access_denied_does_not_fallback_to_standard_key(self): ) self.client._call_flow_frontend_rpc.assert_awaited_once() + async def test_unusual_activity_retries_once_then_raises_and_reports(self): + unusual = RuntimeError( + "Flow frontend RPC rejected: rpc=eb1hJf, error=PUBLIC_ERROR_UNUSUAL_ACTIVITY" + ) + self.client._call_flow_frontend_rpc = AsyncMock(side_effect=unusual) + self.client._notify_browser_captcha_error = AsyncMock() + + with patch("src.services.flow_client.asyncio.sleep", new=AsyncMock()): + with self.assertRaisesRegex(RuntimeError, "UNUSUAL_ACTIVITY"): + await self.client.generate_video_start_image( + at="at-token", + project_id="22222222-2222-2222-2222-222222222222", + prompt="猫猫", + model_key="veo_3_1_i2v_s_fast", + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + start_media_id="33333333-3333-3333-3333-333333333333", + google_cookies="SID=session-cookie", + ) + + # 同账号只重试一次，剩下交给上层换号 + self.assertEqual(self.client._call_flow_frontend_rpc.await_count, 2) + # 每次被拒都通知打码服务自愈 + self.assertEqual(self.client._notify_browser_captcha_error.await_count, 2) + async def test_transient_retry_keeps_original_model_key(self): operation_id = "11111111-1111-1111-1111-111111111111" project_id = "22222222-2222-2222-2222-222222222222" From 52a30362bc612f8d8ae1a2def42d080f2fcb2be2 Mon Sep 17 00:00:00 2001 From: Wizard Date: Wed, 4 Nov 2026 21:44:54 -0400 Subject: [PATCH 2/3] =?UTF-8?q?fix(generation):=20=E9=A2=9D=E5=BA=A6?= =?UTF-8?q?=E4=B8=8D=E8=B6=B3=E8=B4=A6=E5=8F=B7=E4=B8=8D=E5=86=8D=E5=90=8C?= =?UTF-8?q?=E5=8F=B7=E9=87=8D=E8=AF=95=EF=BC=8C=E9=A3=8E=E6=8E=A7=E5=86=B7?= =?UTF-8?q?=E5=8D=B4=E9=80=92=E5=A2=9E=EF=BC=8C=E6=8D=A2=E5=8F=B7=E5=A4=B1?= =?UTF-8?q?=E8=B4=A5=E6=97=B6=E6=B1=87=E6=80=BB=E5=90=84=E8=B4=A6=E5=8F=B7?= =?UTF-8?q?=E5=8E=9F=E5=9B=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - PUBLIC_ERROR_USER_QUOTA_REACHED 不再视为可重试错误，避免在额度不足的账号上 白白消耗多轮打码，直接交给上层换号 - 配额冷却改为最长 6h，并记录耗尽时刷新到的余额；余额回升（充值/额度刷新） 即提前解除，避免低余额账号每小时重新被选中 - 风控冷却按连续被拒次数递增 (5m -> 15m -> 45m -> 60m)，成功一次即重置 - 换号预算拆分：风控失败最多 3 个账号，额度不足的账号不占用该预算， 总尝试上限 6 个账号 - 最终错误信息列出每个失败账号及原因，token_failovers 也记录最后一个账号 Co-Authored-By: Claude Opus 5.5 --- src/services/flow_client.py | 3 ++ src/services/generation_handler.py | 51 ++++++++++++++---- src/services/load_balancer.py | 83 +++++++++++++++++++++++------- tests/test_token_quota_failover.py | 82 +++++++++++++++++++++++++++-- tests/test_veo_lite_support.py | 22 ++++++++ 5 files changed, 211 insertions(+), 30 deletions(-) diff --git a/src/services/flow_client.py b/src/services/flow_client.py index acb03e1f..9aa96b6c 100644 --- a/src/services/flow_client.py +++ b/src/services/flow_client.py @@ -4709,6 +4709,9 @@ def _get_retry_reason(self, error_str: str) -> Optional[str]: if "model_access_denied" in error_lower: # 账号无对应模型权限，重试不会改变结果 return None + if "public_error_user_quota_reached" in error_lower: + # 账号额度不足，同号重试只会白白消耗打码；交给上层换号 + return None if "403" in error_lower: return "403错误" if "429" in error_lower or "too many requests" in error_lower: diff --git a/src/services/generation_handler.py b/src/services/generation_handler.py index ba1cec3d..ef4df86f 100644 --- a/src/services/generation_handler.py +++ b/src/services/generation_handler.py @@ -33,8 +33,11 @@ def _video_poll_attempt_budget( return attempts * (3 if upsample else 1) -# 账号配额耗尽时，单个请求最多尝试的账号数（含首个账号） -MAX_TOKEN_FAILOVER_ATTEMPTS = 3 +# 账号配额耗尽/风控时，单个请求最多尝试的账号数（含首个账号） +MAX_TOKEN_FAILOVER_ATTEMPTS = 6 +# 其中因风控 (UNUSUAL_ACTIVITY) 失败的账号数上限：每次风控要耗费两轮打码， +# 而配额耗尽的账号会立即失败，不占用这部分额度 +MAX_RISK_FAILOVER_ATTEMPTS = 3 # Model configuration @@ -1658,6 +1661,8 @@ async def handle_generation( # 2. 选择Token（账号配额耗尽时自动切换到其他账号重试） failed_token_ids: set = set() last_failover_error: Optional[str] = None + failover_summary: List[str] = [] + risk_failures = 0 for token_attempt in range(MAX_TOKEN_FAILOVER_ATTEMPTS): token = None generation_result = self._create_generation_result() @@ -1702,8 +1707,8 @@ async def handle_generation( error_msg = self._get_no_token_error_message(generation_type) if last_failover_error: error_msg = ( - f"{last_failover_error}（已尝试 {len(failed_token_ids)} 个账号，" - f"没有其他可切换的账号）" + f"{last_failover_error}（已尝试 {len(failed_token_ids)} 个账号均失败: " + f"{', '.join(failover_summary)}；没有其他可切换的账号）" ) debug_logger.log_error(f"[GENERATION] {error_msg}") record_generation_result( @@ -1983,19 +1988,27 @@ async def handle_generation( ): scope = "image" if generation_type == "image" else "video" if is_quota_error: - self.load_balancer.mark_quota_cooldown(token.id, scope) + credits = await self._refresh_token_credits(token) + self.load_balancer.mark_quota_cooldown( + token.id, scope, credits=credits + ) failover_reason = "配额耗尽" failover_notice = "当前账号配额已用尽，正在切换其他账号重试...\n" else: + risk_failures += 1 self.load_balancer.mark_risk_cooldown(token.id, scope) failover_reason = "触发风控" failover_notice = "当前账号触发风控，正在切换其他账号重试...\n" failed_token_ids.add(token.id) last_failover_error = error_msg - if token_attempt < MAX_TOKEN_FAILOVER_ATTEMPTS - 1: - perf_trace.setdefault("token_failovers", []).append( - {"token_id": token.id, "error": str(e)[:200]} - ) + failover_summary.append(f"#{token.id} {failover_reason}") + perf_trace.setdefault("token_failovers", []).append( + {"token_id": token.id, "error": str(e)[:200]} + ) + if ( + token_attempt < MAX_TOKEN_FAILOVER_ATTEMPTS - 1 + and risk_failures < MAX_RISK_FAILOVER_ATTEMPTS + ): debug_logger.log_warning( f"[GENERATION] Token {token.id} {failover_reason}，切换账号重试 " f"({token_attempt + 2}/{MAX_TOKEN_FAILOVER_ATTEMPTS})" @@ -2003,6 +2016,11 @@ async def handle_generation( if stream: yield self._create_stream_chunk(failover_notice) continue + if len(failover_summary) > 1: + error_msg = ( + f"{error_msg}（已尝试 {len(failed_token_ids)} 个账号均失败: " + f"{', '.join(failover_summary)}）" + ) if token: if self._should_count_token_error(e): await self.token_manager.record_error(token.id) @@ -2058,6 +2076,21 @@ def _is_quota_exhausted_error(error: Exception) -> bool: """账号配额耗尽：换一个账号通常可以成功。""" return "public_error_user_quota_reached" in str(error or "").lower() + async def _refresh_token_credits(self, token) -> Optional[int]: + """配额耗尽后刷新并落库余额，作为冷却解除的基准；失败时退回已知余额。""" + try: + return int( + await asyncio.wait_for( + self.token_manager.refresh_credits(token.id), timeout=15 + ) + ) + except Exception as e: + debug_logger.log_warning( + f"[GENERATION] 刷新 Token {token.id} 余额失败: {str(e)[:200]}" + ) + credits = getattr(token, "credits", None) + return credits if isinstance(credits, int) else None + @staticmethod def _is_unusual_activity_error(error: Exception) -> bool: """账号/IP 被风控拒绝：提交被拒、未创建任务，换一个账号通常可以成功。""" diff --git a/src/services/load_balancer.py b/src/services/load_balancer.py index 1f1f0bc1..f6dd0b1a 100644 --- a/src/services/load_balancer.py +++ b/src/services/load_balancer.py @@ -15,10 +15,14 @@ from .concurrency_manager import ConcurrencyManager from ..core.logger import debug_logger -# 配额耗尽 (PUBLIC_ERROR_USER_QUOTA_REACHED) 后降低该 token 优先级的时长 -QUOTA_COOLDOWN_SECONDS = 3600 -# 风控拒绝 (PUBLIC_ERROR_UNUSUAL_ACTIVITY) 后降低该 token 优先级的时长 +# 配额耗尽 (PUBLIC_ERROR_USER_QUOTA_REACHED) 后降低该 token 优先级的最长时长； +# 期间余额回升（充值/额度刷新）会提前解除 +QUOTA_COOLDOWN_SECONDS = 6 * 3600 +# 风控拒绝 (PUBLIC_ERROR_UNUSUAL_ACTIVITY) 后降低该 token 优先级的时长： +# 连续被拒按倍数递增 (5m -> 15m -> 45m -> 60m)，成功一次即重置 RISK_COOLDOWN_SECONDS = 300 +RISK_COOLDOWN_MULTIPLIER = 3 +RISK_COOLDOWN_MAX_SECONDS = 3600 class LoadBalancer: @@ -40,6 +44,10 @@ def __init__( self._rr_lock = asyncio.Lock() # (token_id, "image"/"video") -> monotonic expiry; quota-exhausted tokens sort last self._quota_cooldown: Dict[tuple[int, str], float] = {} + # 配额耗尽时记录的余额；余额高于该值说明已充值，解除冷却 + self._quota_credits: Dict[tuple[int, str], int] = {} + # 连续风控次数，用于递增风控冷却 + self._risk_streak: Dict[tuple[int, str], int] = {} @staticmethod def _scope(for_image_generation: bool, for_video_generation: bool) -> Optional[str]: @@ -50,39 +58,76 @@ def _scope(for_image_generation: bool, for_video_generation: bool) -> Optional[s return None def mark_quota_cooldown( - self, token_id: int, scope: str, seconds: float = QUOTA_COOLDOWN_SECONDS + self, + token_id: int, + scope: str, + seconds: float = QUOTA_COOLDOWN_SECONDS, + credits: Optional[int] = None, ): - """配额耗尽后让该 token 在一段时间内排在候选末尾（不剔除，避免全部冷却时无号可用）。""" - self._quota_cooldown[(token_id, scope)] = time.monotonic() + seconds + """配额耗尽后让该 token 排在候选末尾（不剔除，避免全部冷却时无号可用）。 + + 传入当时的余额后，余额回升即提前解除冷却。 + """ + key = (token_id, scope) + self._quota_cooldown[key] = time.monotonic() + seconds + if credits is None: + self._quota_credits.pop(key, None) + else: + self._quota_credits[key] = credits debug_logger.log_warning( - f"[LOAD_BALANCER] Token {token_id} {scope} 配额耗尽，冷却 {int(seconds)}s 内优先使用其他账号" + f"[LOAD_BALANCER] Token {token_id} {scope} 配额耗尽 (余额: {credits})，" + f"冷却 {int(seconds)}s 内优先使用其他账号" ) def mark_risk_cooldown( - self, token_id: int, scope: str, seconds: float = RISK_COOLDOWN_SECONDS - ): - """风控拒绝后短暂让该 token 排在候选末尾；不会缩短已有的配额冷却。""" + self, token_id: int, scope: str, seconds: Optional[float] = None + ) -> float: + """风控拒绝后让该 token 排在候选末尾，连续被拒时冷却递增；不会缩短已有的配额冷却。""" key = (token_id, scope) + streak = self._risk_streak.get(key, 0) + 1 + self._risk_streak[key] = streak + if seconds is None: + seconds = min( + RISK_COOLDOWN_MAX_SECONDS, + RISK_COOLDOWN_SECONDS * RISK_COOLDOWN_MULTIPLIER ** (streak - 1), + ) expires_at = time.monotonic() + seconds if self._quota_cooldown.get(key, 0) < expires_at: self._quota_cooldown[key] = expires_at debug_logger.log_warning( - f"[LOAD_BALANCER] Token {token_id} {scope} 触发风控，冷却 {int(seconds)}s 内优先使用其他账号" + f"[LOAD_BALANCER] Token {token_id} {scope} 连续第 {streak} 次触发风控，" + f"冷却 {int(seconds)}s 内优先使用其他账号" ) + return seconds def clear_quota_cooldown(self, token_id: int, scope: Optional[str] = None): - for key in list(self._quota_cooldown): - if key[0] == token_id and (scope is None or key[1] == scope): - self._quota_cooldown.pop(key, None) + for state in (self._quota_cooldown, self._quota_credits, self._risk_streak): + for key in list(state): + if key[0] == token_id and (scope is None or key[1] == scope): + state.pop(key, None) - def _in_quota_cooldown(self, token_id: int, scope: Optional[str]) -> bool: + def _in_quota_cooldown( + self, token_id: int, scope: Optional[str], credits: Optional[int] = None + ) -> bool: if scope is None: return False - expires_at = self._quota_cooldown.get((token_id, scope)) + key = (token_id, scope) + expires_at = self._quota_cooldown.get(key) if expires_at is None: return False if time.monotonic() >= expires_at: - self._quota_cooldown.pop((token_id, scope), None) + self._quota_cooldown.pop(key, None) + self._quota_credits.pop(key, None) + return False + exhausted_credits = self._quota_credits.get(key) + if ( + exhausted_credits is not None + and credits is not None + and credits > exhausted_credits + ): + # 余额比耗尽时高，说明已充值或额度已刷新 + self._quota_cooldown.pop(key, None) + self._quota_credits.pop(key, None) return False return True @@ -325,7 +370,9 @@ async def select_token( "inflight": inflight, "remaining": remaining, "needs_refresh": self.token_manager.needs_at_refresh(token), - "quota_cooldown": self._in_quota_cooldown(token.id, scope), + "quota_cooldown": self._in_quota_cooldown( + token.id, scope, getattr(token, "credits", None) + ), "random": random.random(), } ) diff --git a/tests/test_token_quota_failover.py b/tests/test_token_quota_failover.py index 6d6d2706..4bdace54 100644 --- a/tests/test_token_quota_failover.py +++ b/tests/test_token_quota_failover.py @@ -4,10 +4,15 @@ from unittest.mock import AsyncMock, MagicMock, patch from src.services.generation_handler import ( + MAX_RISK_FAILOVER_ATTEMPTS, MAX_TOKEN_FAILOVER_ATTEMPTS, GenerationHandler, ) -from src.services.load_balancer import LoadBalancer +from src.services.load_balancer import ( + RISK_COOLDOWN_MAX_SECONDS, + RISK_COOLDOWN_SECONDS, + LoadBalancer, +) MODEL = "veo-3.1-fast-i2v-6s-landscape" @@ -28,7 +33,7 @@ def _token(token_id): ) -def _build_handler(tokens, quota_token_ids, error=QUOTA_ERROR): +def _build_handler(tokens, quota_token_ids, error=QUOTA_ERROR, errors_by_token=None): token_manager = MagicMock() token_manager.get_active_tokens = AsyncMock(return_value=tokens) token_manager.ensure_valid_token = AsyncMock(side_effect=lambda t: t) @@ -37,6 +42,7 @@ def _build_handler(tokens, quota_token_ids, error=QUOTA_ERROR): token_manager.record_error = AsyncMock() token_manager.record_usage = AsyncMock() token_manager.record_success = AsyncMock() + token_manager.refresh_credits = AsyncMock(return_value=10) load_balancer = LoadBalancer(token_manager) load_balancer._check_extension_route = AsyncMock(return_value=(True, "")) @@ -53,6 +59,8 @@ def _build_handler(tokens, quota_token_ids, error=QUOTA_ERROR): async def fake_video(token, *args, generation_result=None, **kwargs): attempted.append(token.id) + if errors_by_token and token.id in errors_by_token: + raise Exception(errors_by_token[token.id]) if token.id in quota_token_ids: raise error() if callable(error) else Exception(error) handler._mark_generation_succeeded(generation_result) @@ -84,7 +92,7 @@ async def test_switches_to_another_token_on_quota_error(self): self.assertFalse(handler.load_balancer._in_quota_cooldown(2, "video")) async def test_gives_up_after_max_attempts(self): - tokens = [_token(i) for i in range(1, 6)] + tokens = [_token(i) for i in range(1, MAX_TOKEN_FAILOVER_ATTEMPTS + 3)] handler, attempted = _build_handler(tokens, {t.id for t in tokens}) chunks = await _collect(handler) @@ -93,7 +101,23 @@ async def test_gives_up_after_max_attempts(self): self.assertEqual(len(set(attempted)), MAX_TOKEN_FAILOVER_ATTEMPTS) error = json.loads(chunks[-1])["error"] self.assertIn("PUBLIC_ERROR_USER_QUOTA_REACHED", error["message"]) + self.assertIn(f"已尝试 {MAX_TOKEN_FAILOVER_ATTEMPTS} 个账号均失败", error["message"]) self.assertEqual(handler._log_request.await_args.args[4], 500) + # 最后一个失败账号也要记录在性能追踪里 + perf = handler._log_request.await_args.args[3]["performance"] + self.assertEqual(len(perf["token_failovers"]), MAX_TOKEN_FAILOVER_ATTEMPTS) + + async def test_quota_error_records_refreshed_credits(self): + handler, _ = _build_handler([_token(1), _token(2)], {1}) + + with patch("src.services.load_balancer.random.random", side_effect=[0.1, 0.9] * 4): + await _collect(handler) + + handler.token_manager.refresh_credits.assert_awaited_once_with(1) + lb = handler.load_balancer + self.assertTrue(lb._in_quota_cooldown(1, "video", 10)) + # 余额回升（充值/额度刷新）后提前解除冷却 + self.assertFalse(lb._in_quota_cooldown(1, "video", 500)) async def test_reports_quota_error_when_no_token_left(self): handler, attempted = _build_handler([_token(1), _token(2)], {1, 2}) @@ -136,6 +160,58 @@ async def test_risk_cooldown_does_not_shorten_quota_cooldown(self): lb.mark_risk_cooldown(1, "video") self.assertEqual(lb._quota_cooldown[(1, "video")], quota_expiry) + async def test_risk_failures_have_their_own_budget(self): + tokens = [_token(i) for i in range(1, 7)] + handler, attempted = _build_handler( + tokens, set(), errors_by_token={t.id: UNUSUAL_ERROR for t in tokens} + ) + + chunks = await _collect(handler) + + self.assertEqual(len(attempted), MAX_RISK_FAILOVER_ATTEMPTS) + error = json.loads(chunks[-1])["error"]["message"] + self.assertIn("PUBLIC_ERROR_UNUSUAL_ACTIVITY", error) + self.assertIn("触发风控", error) + + async def test_quota_failures_do_not_use_risk_budget(self): + # 复现线上请求：两个账号风控、一个账号额度不足，第四个账号可用 + errors = {7: UNUSUAL_ERROR, 9: UNUSUAL_ERROR, 1: QUOTA_ERROR} + tokens = [_token(i) for i in (7, 9, 1, 5)] + handler, attempted = _build_handler(tokens, set(), errors_by_token=errors) + + # 固定打散顺序为 7 -> 9 -> 1 -> 5（每轮只对未排除的候选取随机数） + order = [0.1, 0.2, 0.3, 0.9, 0.2, 0.3, 0.9, 0.3, 0.9, 0.9] + with patch("src.services.load_balancer.random.random", side_effect=order): + chunks = await _collect(handler) + + self.assertEqual(attempted, [7, 9, 1, 5]) + self.assertNotIn("error", json.loads(chunks[-1])) + + async def test_failure_summary_lists_each_account(self): + errors = {7: UNUSUAL_ERROR, 1: QUOTA_ERROR} + handler, _ = _build_handler( + [_token(7), _token(1)], set(), errors_by_token=errors + ) + + chunks = await _collect(handler) + + error = json.loads(chunks[-1])["error"]["message"] + self.assertIn("已尝试 2 个账号均失败", error) + self.assertIn("#7 触发风控", error) + self.assertIn("#1 配额耗尽", error) + + async def test_risk_cooldown_escalates_and_resets_on_success(self): + handler, _ = _build_handler([_token(1)], set()) + lb = handler.load_balancer + durations = [lb.mark_risk_cooldown(1, "video") for _ in range(5)] + self.assertEqual(durations[0], RISK_COOLDOWN_SECONDS) + self.assertGreater(durations[1], durations[0]) + self.assertEqual(durations[-1], RISK_COOLDOWN_MAX_SECONDS) + + lb.clear_quota_cooldown(1, "video") + self.assertFalse(lb._in_quota_cooldown(1, "video")) + self.assertEqual(lb.mark_risk_cooldown(1, "video"), RISK_COOLDOWN_SECONDS) + async def test_uncertain_submission_is_not_failed_over(self): def uncertain(): error = Exception(UNUSUAL_ERROR) diff --git a/tests/test_veo_lite_support.py b/tests/test_veo_lite_support.py index 99261295..18f0971e 100644 --- a/tests/test_veo_lite_support.py +++ b/tests/test_veo_lite_support.py @@ -361,6 +361,28 @@ async def test_access_denied_does_not_fallback_to_standard_key(self): ) self.client._call_flow_frontend_rpc.assert_awaited_once() + async def test_quota_reached_is_not_retried_on_same_account(self): + self.client._call_flow_frontend_rpc = AsyncMock( + side_effect=RuntimeError( + "Flow frontend RPC rejected: rpc=eb1hJf, error=PUBLIC_ERROR_USER_QUOTA_REACHED" + ) + ) + self.client._notify_browser_captcha_error = AsyncMock() + + with patch("src.services.flow_client.asyncio.sleep", new=AsyncMock()): + with self.assertRaisesRegex(RuntimeError, "QUOTA_REACHED"): + await self.client.generate_video_start_image( + at="at-token", + project_id="22222222-2222-2222-2222-222222222222", + prompt="猫猫", + model_key="veo_3_1_i2v_s_fast", + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + start_media_id="33333333-3333-3333-3333-333333333333", + google_cookies="SID=session-cookie", + ) + + self.client._call_flow_frontend_rpc.assert_awaited_once() + async def test_unusual_activity_retries_once_then_raises_and_reports(self): unusual = RuntimeError( "Flow frontend RPC rejected: rpc=eb1hJf, error=PUBLIC_ERROR_UNUSUAL_ACTIVITY" From 55a7d8c0d777f3b928a2c89bdcfa2d642f6dc456 Mon Sep 17 00:00:00 2001 From: TheSmallHanCat Date: Sun, 8 Nov 2026 13:22:44 +0800 Subject: [PATCH 3/3] fix(generation): isolate failover accounts and preserve cooldown balancing --- src/services/generation_handler.py | 7 +- src/services/load_balancer.py | 7 +- src/services/token_manager.py | 11 ++- tests/test_token_quota_failover.py | 146 ++++++++++++++++++++++++++++- 4 files changed, 166 insertions(+), 5 deletions(-) diff --git a/src/services/generation_handler.py b/src/services/generation_handler.py index d351e4c1..1f2fee54 100644 --- a/src/services/generation_handler.py +++ b/src/services/generation_handler.py @@ -1728,6 +1728,10 @@ async def handle_generation( token = None generation_result = self._create_generation_result() pending_token_state = {"active": False} + if token_attempt > 0 and hasattr( + self.flow_client, "clear_request_fingerprint" + ): + self.flow_client.clear_request_fingerprint() # 选择本次尝试的 Token debug_logger.log_info(f"[GENERATION] 正在选择可用Token...") token_select_started_at = time.time() @@ -2142,7 +2146,8 @@ async def _refresh_token_credits(self, token) -> Optional[int]: try: return int( await asyncio.wait_for( - self.token_manager.refresh_credits(token.id), timeout=15 + self.token_manager.refresh_credits(token.id, raise_on_error=True), + timeout=15, ) ) except Exception as e: diff --git a/src/services/load_balancer.py b/src/services/load_balancer.py index f6dd0b1a..ad67db73 100644 --- a/src/services/load_balancer.py +++ b/src/services/load_balancer.py @@ -398,7 +398,12 @@ async def select_token( scenario = "video" ordered_candidates = [] - first_candidate = await self._select_round_robin(available_tokens, scenario) + round_robin_candidates = [ + item for item in available_tokens if not item["quota_cooldown"] + ] or available_tokens + first_candidate = await self._select_round_robin( + round_robin_candidates, scenario + ) if first_candidate is not None: ordered_candidates.append(first_candidate) ordered_candidates.extend( diff --git a/src/services/token_manager.py b/src/services/token_manager.py index a8bca01c..64829997 100644 --- a/src/services/token_manager.py +++ b/src/services/token_manager.py @@ -1437,19 +1437,26 @@ async def auto_unban_429_tokens(self): # ========== 余额刷新 ========== - async def refresh_credits(self, token_id: int) -> int: + async def refresh_credits(self, token_id: int, *, raise_on_error: bool = False) -> int: """刷新Token余额 + Args: + raise_on_error: 刷新失败时抛出异常，避免将失败误认为真实零余额。 + Returns: credits """ token = await self.db.get_token(token_id) if not token: + if raise_on_error: + raise ValueError("Token not found") return 0 # 确保AT有效 token = await self.ensure_valid_token(token) if not token: + if raise_on_error: + raise ValueError("Token validation failed") return 0 try: @@ -1473,4 +1480,6 @@ async def refresh_credits(self, token_id: int) -> int: debug_logger.log_error( f"Failed to refresh credits for token {token_id}: {str(e)}" ) + if raise_on_error: + raise return 0 diff --git a/tests/test_token_quota_failover.py b/tests/test_token_quota_failover.py index 4bdace54..81cbaaeb 100644 --- a/tests/test_token_quota_failover.py +++ b/tests/test_token_quota_failover.py @@ -1,7 +1,10 @@ import json import unittest from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +from src.core.config import config +from src.services.flow_client import FlowClient from src.services.generation_handler import ( MAX_RISK_FAILOVER_ATTEMPTS, @@ -13,6 +16,7 @@ RISK_COOLDOWN_SECONDS, LoadBalancer, ) +from src.services.token_manager import TokenManager MODEL = "veo-3.1-fast-i2v-6s-landscape" @@ -113,7 +117,9 @@ async def test_quota_error_records_refreshed_credits(self): with patch("src.services.load_balancer.random.random", side_effect=[0.1, 0.9] * 4): await _collect(handler) - handler.token_manager.refresh_credits.assert_awaited_once_with(1) + handler.token_manager.refresh_credits.assert_awaited_once_with( + 1, raise_on_error=True + ) lb = handler.load_balancer self.assertTrue(lb._in_quota_cooldown(1, "video", 10)) # 余额回升（充值/额度刷新）后提前解除冷却 @@ -226,5 +232,141 @@ def uncertain(): self.assertIn("PUBLIC_ERROR_UNUSUAL_ACTIVITY", json.loads(chunks[-1])["error"]["message"]) +class TokenFailoverIsolationTests(unittest.IsolatedAsyncioTestCase): + async def test_next_account_prepares_project_without_previous_fingerprint(self): + for error in (QUOTA_ERROR, UNUSUAL_ERROR): + with self.subTest(error=error): + handler, attempted = _build_handler( + [_token(1), _token(2)], {1}, error + ) + flow_client = FlowClient(proxy_manager=None) + flow_client.prefill_remote_browser_pool = AsyncMock() + handler.flow_client = flow_client + project_fingerprints = {} + + async def prepare_project(token_id): + project_fingerprints[token_id] = ( + flow_client.get_request_fingerprint() + ) + return f"project-{token_id}" + + generate_video = handler._handle_video_generation + + async def video_with_fingerprint(token, *args, **kwargs): + flow_client._set_request_fingerprint( + { + "session_cookies": {"OSID": f"session-{token.id}"}, + "proxy_url": f"http://proxy-{token.id}:8080", + } + ) + async for chunk in generate_video(token, *args, **kwargs): + yield chunk + + handler.token_manager.ensure_project_exists = AsyncMock( + side_effect=prepare_project + ) + handler._handle_video_generation = video_with_fingerprint + with patch( + "src.services.load_balancer.random.random", + side_effect=[0.1, 0.9] * 4, + ): + chunks = await _collect(handler) + + self.assertEqual(attempted, [1, 2]) + self.assertNotIn("error", json.loads(chunks[-1])) + self.assertEqual(project_fingerprints, {1: None, 2: None}) + + +class TokenCreditRefreshTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.token = _token(1) + self.token.at = "unused" + self.token.google_cookies = "" + self.manager = TokenManager.__new__(TokenManager) + self.manager.db = MagicMock() + self.manager.db.get_token = AsyncMock(return_value=self.token) + self.manager.db.update_token = AsyncMock() + self.manager.ensure_valid_token = AsyncMock(return_value=self.token) + self.manager.flow_client = MagicMock() + self.manager.flow_client.get_credits = AsyncMock() + self.handler = GenerationHandler.__new__(GenerationHandler) + self.handler.token_manager = self.manager + + async def test_refresh_failure_keeps_cooldown_with_cached_balance(self): + self.manager.flow_client.get_credits.side_effect = OSError("connection failed") + + credits = await self.handler._refresh_token_credits(self.token) + lb = LoadBalancer(self.manager) + lb.mark_quota_cooldown(self.token.id, "video", credits=credits) + + self.assertEqual(credits, self.token.credits) + self.manager.db.update_token.assert_not_awaited() + self.assertTrue( + lb._in_quota_cooldown(self.token.id, "video", self.token.credits) + ) + + async def test_successful_refresh_preserves_real_zero_balance(self): + self.manager.flow_client.get_credits.return_value = { + "credits": 0, + "userPaygateTier": "PAYGATE_TIER_ONE", + } + + credits = await self.handler._refresh_token_credits(self.token) + + self.assertEqual(credits, 0) + self.manager.db.update_token.assert_awaited_once_with( + self.token.id, credits=0, user_paygate_tier="PAYGATE_TIER_ONE" + ) + + async def test_default_refresh_keeps_existing_failure_contract(self): + self.manager.flow_client.get_credits.side_effect = OSError("connection failed") + + self.assertEqual(await self.manager.refresh_credits(self.token.id), 0) + self.manager.db.update_token.assert_not_awaited() + + +class TokenCooldownPollingTests(unittest.IsolatedAsyncioTestCase): + async def test_polling_rotates_between_accounts_outside_cooldown(self): + for cooldown_method in ("mark_quota_cooldown", "mark_risk_cooldown"): + with self.subTest(cooldown_method=cooldown_method): + handler, _ = _build_handler([_token(i) for i in range(1, 5)], set()) + lb = handler.load_balancer + for token_id in (2, 3): + getattr(lb, cooldown_method)(token_id, "video") + lb._round_robin_state["video"] = 1 + + with patch.object( + type(config), + "call_logic_mode", + new_callable=PropertyMock, + return_value="polling", + ): + selected = [ + (await lb.select_token(for_video_generation=True)).id + for _ in range(4) + ] + + self.assertEqual(selected, [4, 1, 4, 1]) + + async def test_polling_keeps_cooled_accounts_as_fallback(self): + handler, _ = _build_handler([_token(1), _token(2)], set()) + lb = handler.load_balancer + for token_id in (1, 2): + lb.mark_quota_cooldown(token_id, "video") + + with patch.object( + type(config), + "call_logic_mode", + new_callable=PropertyMock, + return_value="polling", + ): + selected = [ + (await lb.select_token(for_video_generation=True)).id + for _ in range(4) + ] + + self.assertEqual(selected, [1, 2, 1, 2]) + + if __name__ == "__main__": unittest.main()