diff --git a/cat.jpg b/cat.jpg deleted file mode 100644 index 08f4891..0000000 Binary files a/cat.jpg and /dev/null differ diff --git a/cat_watercolor.jpg b/cat_watercolor.jpg deleted file mode 100644 index 2b6697c..0000000 Binary files a/cat_watercolor.jpg and /dev/null differ diff --git a/docker-compose.yml b/docker-compose.yml index d02d23b..ae005d0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -4,14 +4,12 @@ services: flow2api: image: ghcr.io/thesmallhancat/flow2api:latest container_name: flow2api - network_mode: host + ports: + - "38000:8000" volumes: - ./data:/app/data - ./tmp:/app/tmp - ./config/setting.toml:/app/config/setting.toml - - ./src/services/proxy_manager.py:/app/src/services/proxy_manager.py - - ./src/services/flow_client.py:/app/src/services/flow_client.py - - ./src/api/admin.py:/app/src/api/admin.py environment: - PYTHONUNBUFFERED=1 restart: unless-stopped diff --git a/extension/background.js b/extension/background.js index 22ec2cc..517e3f6 100644 --- a/extension/background.js +++ b/extension/background.js @@ -4,6 +4,7 @@ let heartbeatInterval = null; const DEFAULT_SETTINGS = { serverUrl: "ws://127.0.0.1:8000/captcha_ws", + apiKey: "", routeKey: "", clientLabel: "" }; @@ -13,6 +14,7 @@ function getSettings() { chrome.storage.local.get(DEFAULT_SETTINGS, (stored) => { resolve({ serverUrl: (stored.serverUrl || DEFAULT_SETTINGS.serverUrl).trim(), + apiKey: (stored.apiKey || "").trim(), routeKey: (stored.routeKey || "").trim(), clientLabel: (stored.clientLabel || "").trim() }); @@ -74,6 +76,9 @@ async function connectWS() { const settings = await getSettings(); const url = new URL(settings.serverUrl || DEFAULT_SETTINGS.serverUrl); + if (settings.apiKey) { + url.searchParams.set("key", settings.apiKey); + } if (settings.routeKey) { url.searchParams.set("route_key", settings.routeKey); } @@ -227,7 +232,7 @@ async function handleGetToken(data) { chrome.storage.onChanged.addListener((changes, areaName) => { if (areaName !== "local") return; - if (changes.routeKey || changes.serverUrl || changes.clientLabel) { + if (changes.routeKey || changes.serverUrl || changes.apiKey || changes.clientLabel) { console.log("[Flow2API] Extension settings changed, reconnecting WebSocket..."); closeSocket(); connectWS(); diff --git a/extension/manifest 2.json b/extension/manifest 2.json deleted file mode 100644 index 7da2a03..0000000 --- a/extension/manifest 2.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "manifest_version": 3, - "name": "Flow2API Captcha Worker", - "version": "1.0", - "permissions": ["activeTab", "scripting", "tabs", "storage"], - "host_permissions": ["https://labs.google/*"], - "options_page": "options.html", - "background": { - "service_worker": "background.js" - }, - "content_scripts": [ - { - "matches": ["https://labs.google/*"], - "js": ["content.js"], - "run_at": "document_end" - } - ] -} diff --git a/extension/manifest.json b/extension/manifest.json index fa2b223..f69c33e 100644 --- a/extension/manifest.json +++ b/extension/manifest.json @@ -9,7 +9,7 @@ "scripting" ], "host_permissions": [ - "" + "https://labs.google/*" ], "background": { "service_worker": "background.js" diff --git a/extension/options 2.js b/extension/options 2.js deleted file mode 100644 index e534215..0000000 --- a/extension/options 2.js +++ /dev/null @@ -1,47 +0,0 @@ -const DEFAULTS = { - routeKey: "", - clientLabel: "", - serverUrl: "ws://127.0.0.1:8000/captcha_ws" -}; - -function $(id) { - return document.getElementById(id); -} - -function showStatus(message, isError = false) { - const status = $("status"); - status.textContent = message; - status.style.color = isError ? "#b91c1c" : "#065f46"; -} - -function loadSettings() { - chrome.storage.local.get(DEFAULTS, (stored) => { - $("routeKey").value = stored.routeKey || ""; - $("clientLabel").value = stored.clientLabel || ""; - $("serverUrl").value = stored.serverUrl || DEFAULTS.serverUrl; - }); -} - -function saveSettings() { - const routeKey = $("routeKey").value.trim(); - const clientLabel = $("clientLabel").value.trim(); - const serverUrl = $("serverUrl").value.trim() || DEFAULTS.serverUrl; - - if (!serverUrl.startsWith("ws://") && !serverUrl.startsWith("wss://")) { - showStatus("WebSocket URL 必须以 ws:// 或 wss:// 开头", true); - return; - } - - chrome.storage.local.set({ routeKey, clientLabel, serverUrl }, () => { - if (chrome.runtime.lastError) { - showStatus(chrome.runtime.lastError.message || "保存失败", true); - return; - } - showStatus("已保存，后台连接会自动重连"); - }); -} - -document.addEventListener("DOMContentLoaded", () => { - loadSettings(); - $("saveBtn").addEventListener("click", saveSettings); -}); diff --git a/extension/options.html b/extension/options.html index 99f79ee..4bded69 100644 --- a/extension/options.html +++ b/extension/options.html @@ -87,9 +87,12 @@ + + +
-
管理台 token 里填写同样的 Route Key，比如 9223 对 9223。
+
管理台 token 里填写同样的 Route Key，比如 9223 对 9223。后端会使用 API Key 校验 WebSocket 连接。
diff --git a/extension/options.js b/extension/options.js index bb00375..09ebcdb 100644 --- a/extension/options.js +++ b/extension/options.js @@ -1,5 +1,6 @@ const DEFAULT_SETTINGS = { serverUrl: "ws://127.0.0.1:8000/captcha_ws", + apiKey: "", routeKey: "", clientLabel: "" }; @@ -9,6 +10,7 @@ const $ = (id) => document.getElementById(id); function normalizeSettings(values) { return { serverUrl: (values.serverUrl || DEFAULT_SETTINGS.serverUrl).trim(), + apiKey: (values.apiKey || "").trim(), routeKey: (values.routeKey || "").trim(), clientLabel: (values.clientLabel || "").trim() }; @@ -33,6 +35,7 @@ function loadSettings() { chrome.storage.local.get(DEFAULT_SETTINGS, (stored) => { const settings = normalizeSettings(stored); $("serverUrl").value = settings.serverUrl; + $("apiKey").value = settings.apiKey; $("routeKey").value = settings.routeKey; $("clientLabel").value = settings.clientLabel; }); @@ -41,6 +44,7 @@ function loadSettings() { function saveSettings() { const settings = normalizeSettings({ serverUrl: $("serverUrl").value, + apiKey: $("apiKey").value, routeKey: $("routeKey").value, clientLabel: $("clientLabel").value }); @@ -49,6 +53,10 @@ function saveSettings() { setStatus("WebSocket URL 必须以 ws:// 或 wss:// 开头。", true); return; } + if (!settings.apiKey) { + setStatus("请填写 Flow2API API Key。", true); + return; + } chrome.storage.local.set(settings, () => { if (chrome.runtime.lastError) { diff --git a/src/api/routes.py b/src/api/routes.py index a7caa4b..c843a59 100644 --- a/src/api/routes.py +++ b/src/api/routes.py @@ -9,10 +9,10 @@ import re from urllib.parse import urlparse from curl_cffi.requests import AsyncSession -from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket, WebSocketDisconnect +from fastapi import APIRouter, Depends, HTTPException, Query, Request, WebSocket, WebSocketDisconnect from fastapi.responses import JSONResponse, StreamingResponse -from ..core.auth import verify_api_key_flexible +from ..core.auth import AuthManager, verify_api_key_flexible from ..core.logger import debug_logger from ..core.model_resolver import get_base_model_aliases, resolve_model_name from ..core.models import ( @@ -23,7 +23,6 @@ from ..core.models import ( ) from ..services.generation_handler import MODEL_CONFIG, GenerationHandler from ..services.browser_captcha_extension import ExtensionCaptchaService -from fastapi import WebSocket, WebSocketDisconnect router = APIRouter() @@ -486,6 +485,7 @@ async def _collect_non_stream_result( model: str, prompt: str, images: List[bytes], + base_url_override: Optional[str] = None, video_media_id: Optional[str] = None, ) -> str: handler = _ensure_generation_handler() @@ -495,6 +495,7 @@ async def _collect_non_stream_result( prompt=prompt, images=images if images else None, stream=False, + base_url_override=base_url_override, video_media_id=video_media_id, ): result = chunk @@ -723,6 +724,7 @@ async def _iterate_openai_stream( prompt=normalized.prompt, images=normalized.images if normalized.images else None, stream=True, + base_url_override=base_url_override, video_media_id=normalized.video_media_id, ): if chunk.startswith("data: "): @@ -746,6 +748,7 @@ async def _iterate_gemini_stream( prompt=normalized.prompt, images=normalized.images if normalized.images else None, stream=True, + base_url_override=base_url_override, video_media_id=normalized.video_media_id, ): if chunk.startswith("data: "): @@ -874,6 +877,7 @@ async def create_chat_completion( normalized.model, normalized.prompt, normalized.images, + base_url_override=request_base_url, video_media_id=normalized.video_media_id, ) ) @@ -907,7 +911,8 @@ async def generate_content( normalized.model, normalized.prompt, normalized.images, - request_base_url, + base_url_override=request_base_url, + video_media_id=normalized.video_media_id, ) ) ) @@ -970,6 +975,20 @@ async def stream_generate_content( @router.websocket("/captcha_ws") async def captcha_websocket_endpoint(websocket: WebSocket): from ..core.logger import debug_logger + api_key = ( + websocket.query_params.get("key") + or websocket.query_params.get("api_key") + or websocket.headers.get("x-goog-api-key") + or "" + ).strip() + authorization = (websocket.headers.get("authorization") or "").strip() + if authorization.lower().startswith("bearer "): + api_key = authorization[7:].strip() + + if not api_key or not AuthManager.verify_api_key(api_key): + await websocket.close(code=1008) + return + service = await ExtensionCaptchaService.get_instance() await service.connect(websocket) try: diff --git a/src/core/config.py b/src/core/config.py index 10f3449..ab3b968 100644 --- a/src/core/config.py +++ b/src/core/config.py @@ -12,8 +12,11 @@ class Config: self._admin_password: Optional[str] = None def _load_config(self) -> Dict[str, Any]: - """Load configuration from setting.toml""" - config_path = Path(__file__).parent.parent.parent / "config" / "setting.toml" + """Load configuration from setting.toml, falling back to the example file.""" + config_dir = Path(__file__).parent.parent.parent / "config" + config_path = config_dir / "setting.toml" + if not config_path.exists(): + config_path = config_dir / "setting_example.toml" with open(config_path, "rb") as f: return tomli.load(f) diff --git a/src/core/database.py b/src/core/database.py index 02a1e1d..be06852 100644 --- a/src/core/database.py +++ b/src/core/database.py @@ -408,6 +408,7 @@ class Database: ("image_concurrency", "INTEGER DEFAULT -1"), ("video_concurrency", "INTEGER DEFAULT -1"), ("captcha_proxy_url", "TEXT"), # token级打码代理 + ("extension_route_key", "TEXT"), # extension 模式路由键 ("ban_reason", "TEXT"), # 禁用原因 ("banned_at", "TIMESTAMP"), # 禁用时间 ] @@ -567,6 +568,7 @@ class Database: image_concurrency INTEGER DEFAULT -1, video_concurrency INTEGER DEFAULT -1, captcha_proxy_url TEXT, + extension_route_key TEXT, ban_reason TEXT, banned_at TIMESTAMP ) @@ -845,13 +847,15 @@ class Database: cursor = await db.execute(""" INSERT INTO tokens (st, at, at_expires, email, name, remark, is_active, credits, user_paygate_tier, current_project_id, current_project_name, - image_enabled, video_enabled, image_concurrency, video_concurrency, captcha_proxy_url) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + image_enabled, video_enabled, image_concurrency, video_concurrency, + captcha_proxy_url, extension_route_key) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, (token.st, token.at, token.at_expires, token.email, token.name, token.remark, token.is_active, token.credits, token.user_paygate_tier, token.current_project_id, token.current_project_name, token.image_enabled, token.video_enabled, - token.image_concurrency, token.video_concurrency, token.captcha_proxy_url)) + token.image_concurrency, token.video_concurrency, + token.captcha_proxy_url, token.extension_route_key)) await db.commit() token_id = cursor.lastrowid diff --git a/src/core/models.py b/src/core/models.py index c9f1d37..df65b0b 100644 --- a/src/core/models.py +++ b/src/core/models.py @@ -42,6 +42,7 @@ class Token(BaseModel): # 打码代理（token 级，可覆盖全局浏览器打码代理） captcha_proxy_url: Optional[str] = None + extension_route_key: Optional[str] = None # 429禁用相关 ban_reason: Optional[str] = None # 禁用原因: "429_rate_limit" 或 None diff --git a/src/services/browser_captcha_extension.py b/src/services/browser_captcha_extension.py index 107e367..e323a9e 100644 --- a/src/services/browser_captcha_extension.py +++ b/src/services/browser_captcha_extension.py @@ -1,6 +1,7 @@ import asyncio import json import time +import uuid from dataclasses import dataclass, field from typing import Any, Dict, Optional @@ -24,7 +25,7 @@ class ExtensionCaptchaService: def __init__(self, db=None): self.db = db self.active_connections: list[ExtensionConnection] = [] - self.pending_requests: dict[str, asyncio.Future] = {} + self.pending_requests: dict[str, tuple[asyncio.Future, WebSocket]] = {} @classmethod async def get_instance(cls, db=None) -> "ExtensionCaptchaService": @@ -143,7 +144,10 @@ class ExtensionCaptchaService: req_id = payload.get("req_id") if req_id and req_id in self.pending_requests: - future = self.pending_requests[req_id] + future, owner_websocket = self.pending_requests[req_id] + if websocket is not owner_websocket: + debug_logger.log_warning(f"[Extension Captcha] Ignoring response from non-owner connection: {req_id}") + return if not future.done(): future.set_result(payload) except Exception as e: @@ -169,9 +173,9 @@ class ExtensionCaptchaService: f"Available route keys: {available}" ) - req_id = f"req_{int(time.time() * 1000)}_{id(self)}" - future = asyncio.Future() - self.pending_requests[req_id] = future + req_id = f"req_{uuid.uuid4().hex}" + future = asyncio.get_running_loop().create_future() + self.pending_requests[req_id] = (future, conn.websocket) request_data = { "type": "get_token", @@ -194,16 +198,13 @@ class ExtensionCaptchaService: error_msg = result.get("error") debug_logger.log_error(f"[Extension Captcha] Error from extension: {error_msg}") - print(f"!!! EXTENSION ERROR: {error_msg} !!!") return None except asyncio.TimeoutError: debug_logger.log_error(f"[Extension Captcha] Timeout waiting for token (req_id: {req_id})") - print("!!! EXTENSION TIMEOUT !!!") return None except Exception as e: debug_logger.log_error(f"[Extension Captcha] Communication error: {e}") - print(f"!!! EXTENSION COMM ERROR: {e} !!!") return None finally: self.pending_requests.pop(req_id, None) diff --git a/src/services/flow_client.py b/src/services/flow_client.py index bf39d5a..c4e66c7 100644 --- a/src/services/flow_client.py +++ b/src/services/flow_client.py @@ -213,14 +213,6 @@ class FlowClient: headers["sec-ch-ua-platform"] = "\"Windows\"" headers["sec-ch-ua-mobile"] = "?0" - if 'aisandbox' in url: - print(f"[DEBUG-DEEP] API REQUEST to: {url[:80]}") - print(f"[DEBUG-DEEP] fingerprint: {fingerprint}") - print(f"[DEBUG-DEEP] proxy_url: {proxy_url}") - print(f"[DEBUG-DEEP] UA: {headers.get('User-Agent', '')[:100]}") - print(f"[DEBUG-DEEP] sec-ch-ua-platform: {headers.get('sec-ch-ua-platform', 'NOT SET')}") - print(f"[DEBUG-DEEP] sec-ch-ua-mobile: {headers.get('sec-ch-ua-mobile', 'NOT SET')}") - # Log request if config.debug_enabled: if isinstance(fingerprint, dict): @@ -1946,7 +1938,6 @@ class FlowClient: ] } - print(f"[CONCAT] 提交拼接: original={original_media_id}, extend={extend_media_id}", flush=True) debug_logger.log_info(f"[CONCAT] 提交拼接任务: original={original_media_id[:12]}..., extend={extend_media_id[:12]}...") result = await self._make_request( @@ -1956,7 +1947,7 @@ class FlowClient: use_at=True, at_token=at ) - print(f"[CONCAT] 提交结果: {json.dumps(result, ensure_ascii=False)[:500]}", flush=True) + debug_logger.log_info(f"[CONCAT] 拼接任务已提交: {json.dumps(result, ensure_ascii=False)[:300]}") return result async def poll_concatenation_status( @@ -2006,11 +1997,14 @@ class FlowClient: ev_len = len(encoded_video) if encoded_video else 0 elapsed = int(time.time() - start_time) all_keys = list(result.keys()) - print(f"[CONCAT] 状态: {status}, outputUri: {'有:'+output_uri[:40] if output_uri else 'N/A'}, encodedVideo: {ev_len} chars, elapsed: {elapsed}s, keys: {all_keys}", flush=True) + debug_logger.log_info( + f"[CONCAT] 状态: {status}, outputUri={'yes' if output_uri else 'no'}, " + f"encodedVideo={ev_len} chars, elapsed={elapsed}s, keys={all_keys}" + ) # 优先检查 outputUri if output_uri: - print(f"[CONCAT] 拼接完成 (outputUri): {output_uri[:120]}", flush=True) + debug_logger.log_info(f"[CONCAT] 拼接完成 (outputUri): {output_uri[:120]}") return result # Google API 返回 encodedVideo（base64 编码的 MP4）而不是 outputUri @@ -2030,26 +2024,26 @@ class FlowClient: # 构造 URL：FastAPI 挂载了 /tmp -> /app/tmp/ serve_url = f"/tmp/{video_filename}" - print(f"[CONCAT] 拼接完成 (encodedVideo): 保存 {len(video_bytes)} bytes -> {serve_url}", flush=True) + debug_logger.log_info(f"[CONCAT] 拼接完成 (encodedVideo): 保存 {len(video_bytes)} bytes -> {serve_url}") result["outputUri"] = serve_url result["local_file"] = save_path return result except Exception as e: - print(f"[CONCAT] 解码 encodedVideo 失败: {e}", flush=True) + debug_logger.log_error(f"[CONCAT] 解码 encodedVideo 失败: {e}") raise Exception(f"解码拼接视频失败: {e}") # SUCCESSFUL but neither outputUri nor encodedVideo if "SUCCESSFUL" in status: - print(f"[CONCAT] SUCCESSFUL 但无 outputUri/encodedVideo，完整响应: {json.dumps(result, ensure_ascii=False)[:500]}", flush=True) - + debug_logger.log_warning(f"[CONCAT] SUCCESSFUL 但无 outputUri/encodedVideo: {json.dumps(result, ensure_ascii=False)[:300]}") + if "FAILED" in status or "ERROR" in status: - print(f"[CONCAT] 失败: {status}, 响应: {json.dumps(result, ensure_ascii=False)[:300]}", flush=True) + debug_logger.log_error(f"[CONCAT] 失败: {status}, 响应: {json.dumps(result, ensure_ascii=False)[:300]}") raise Exception(f"视频拼接失败: {status}") await asyncio.sleep(poll_interval) - print(f"[CONCAT] 超时 ({timeout}s)，放弃拼接", flush=True) + debug_logger.log_error(f"[CONCAT] 超时 ({timeout}s)，放弃拼接") raise Exception(f"视频拼接超时 ({timeout}s)") # ========== 视频放大 (Video Upsampler) ========== @@ -2662,13 +2656,12 @@ class FlowClient: - 其他模式: browser_id 为 None """ captcha_method = config.captcha_method - print(f"[DEBUG] _get_recaptcha_token called: captcha_method={captcha_method}, project_id={project_id}, action={action}") + debug_logger.log_info(f"[reCAPTCHA] 开始获取 token: method={captcha_method}, project_id={project_id}, action={action}") if captcha_method == "extension": try: from .browser_captcha_extension import ExtensionCaptchaService service = await ExtensionCaptchaService.get_instance(self.db) - print(f"[DEBUG-EXTENSION] get_token requested, active_connections length: {len(service.active_connections)}") extension_timeout = 45 if action == "VIDEO_GENERATION" else 25 token = await service.get_token( project_id, @@ -2680,7 +2673,6 @@ class FlowClient: return token, None except Exception as e: debug_logger.log_error(f"[reCAPTCHA Extension] 错误: {str(e)}") - print(f"[DEBUG-EXTENSION] EXCEPTION IN get_token: {e}") self._set_request_fingerprint(None) return None, None @@ -2695,8 +2687,6 @@ class FlowClient: token = await service.get_token(project_id, action) debug_logger.log_info(f"[reCAPTCHA] get_token 返回: {token[:50] if token else None}...") fingerprint = service.get_last_fingerprint() if token else None - print(f"[DEBUG-DEEP] personal token obtained: {bool(token)}, token_prefix={str(token)[:40] if token else 'None'}") - print(f"[DEBUG-DEEP] personal fingerprint: {fingerprint}") self._set_request_fingerprint(fingerprint if token else None) return token, None except RuntimeError as e: @@ -2772,11 +2762,10 @@ class FlowClient: proxy_url = await self.proxy_manager.get_request_proxy_url() if proxy_url: self._set_request_fingerprint({"proxy_url": proxy_url}) - print(f"[DEBUG] API captcha using proxy: {proxy_url[:60]}...") else: self._set_request_fingerprint(None) except Exception as e: - print(f"[DEBUG] Failed to get proxy for API captcha: {e}") + debug_logger.log_warning(f"[reCAPTCHA] Failed to get proxy for API captcha: {e}") self._set_request_fingerprint(None) else: self._set_request_fingerprint(None) @@ -2832,22 +2821,15 @@ class FlowClient: if self.proxy_manager: try: proxy_url = await self.proxy_manager.get_request_proxy_url() - print(f"[DEBUG] Got proxy_url: {proxy_url[:60] if proxy_url else None}") if proxy_url: if proxy_url.startswith("socks5://"): # curl_cffi 对 SOCKS5 使用 proxy 参数 proxy = proxy_url - print(f"[DEBUG] Using SOCKS5 proxy for CapSolver: {proxy[:60]}...") else: # HTTP/HTTPS 代理使用 proxies 字典 proxies = {"http": proxy_url, "https": proxy_url} - print(f"[DEBUG] Using HTTP proxy for CapSolver: {proxy_url[:60]}...") except Exception as e: - print(f"[DEBUG] Failed to get proxy: {e}") debug_logger.log_warning(f"[reCAPTCHA {method}] Failed to get proxy: {e}") - - print(f"[DEBUG] CapSolver request: method={method}, project_id={project_id}, action={action}") - print(f"[DEBUG] CapSolver API Key: {client_key[:20]}...") async with AsyncSession() as session: create_url = f"{base_url}/createTask" @@ -2861,21 +2843,15 @@ class FlowClient: } } - # 根据代理类型使用不同参数 - print(f"[DEBUG] Sending createTask request to {create_url}") if proxy: - print(f"[DEBUG] Using proxy param: {proxy[:60]}...") result = await session.post(create_url, json=create_data, impersonate="chrome124", proxy=proxy) else: - print(f"[DEBUG] Using proxies dict or no proxy") result = await session.post(create_url, json=create_data, impersonate="chrome124", proxies=proxies) - - print(f"[DEBUG] createTask response status: {result.status_code}") + debug_logger.log_info(f"[reCAPTCHA {method}] createTask response status: {result.status_code}") result_json = result.json() task_id = result_json.get('taskId') - print(f"[DEBUG] createTask response: {result_json}") debug_logger.log_info(f"[reCAPTCHA {method}] created task_id: {task_id}, response: {result_json}") if not task_id: @@ -2896,29 +2872,21 @@ class FlowClient: result = await session.post(get_url, json=get_data, impersonate="chrome124", proxies=proxies) result_json = result.json() - print(f"[DEBUG] getTaskResult polling #{i+1}: status={result_json.get('status')}, response={result_json}") debug_logger.log_info(f"[reCAPTCHA {method}] polling #{i+1}: {result_json}") status = result_json.get('status') if status == 'ready': solution = result_json.get('solution', {}) response = solution.get('gRecaptchaResponse') - print(f"[DEBUG] Got solution: {solution}") - print(f"[DEBUG] Got gRecaptchaResponse: {response[:50] if response else None}") if response: - print(f"[DEBUG] Token获取成功") debug_logger.log_info(f"[reCAPTCHA {method}] Token获取成功") return response await asyncio.sleep(3) - print(f"[DEBUG] Timeout waiting for token after 40 retries") debug_logger.log_error(f"[reCAPTCHA {method}] Timeout waiting for token") return None except Exception as e: - print(f"[DEBUG] Exception in _get_api_captcha_token: {type(e).__name__}: {str(e)}") debug_logger.log_error(f"[reCAPTCHA {method}] error: {str(e)}") - import traceback - traceback.print_exc() return None diff --git a/src/services/generation_handler.py b/src/services/generation_handler.py index b7cd215..243f668 100644 --- a/src/services/generation_handler.py +++ b/src/services/generation_handler.py @@ -818,6 +818,7 @@ class GenerationHandler: prompt: str, images: Optional[List[bytes]] = None, stream: bool = False, + base_url_override: Optional[str] = None, video_media_id: Optional[str] = None, ) -> AsyncGenerator: """统一生成入口 @@ -1688,7 +1689,6 @@ class GenerationHandler: debug_logger.log_info(f"[EXTEND] 续写视频: {video_media_id}") if stream: yield self._create_stream_chunk(f"视频续写任务提交中，源视频: {video_media_id[:8]}...\n") - print(f"[EXTEND-DEBUG] Calling generate_video_extend with video_media_id={video_media_id}") result = await self.flow_client.generate_video_extend( at=token.at, project_id=project_id, @@ -1755,7 +1755,17 @@ class GenerationHandler: # 如果是 extend，传入源视频 media_id 用于后续拼接 extend_source_id = video_media_id if video_type == "extend" else None - async for chunk in self._poll_video_result(token, project_id, operations, stream, upsample_config, generation_result, request_log_state, extend_source_media_id=extend_source_id): + async for chunk in self._poll_video_result( + token, + project_id, + operations, + stream, + upsample_config, + generation_result, + response_state, + request_log_state, + extend_source_media_id=extend_source_id, + ): yield chunk finally: @@ -1769,6 +1779,7 @@ class GenerationHandler: stream: bool, upsample_config: Optional[Dict] = None, generation_result: Optional[Dict[str, Any]] = None, + response_state: Optional[Dict[str, Any]] = None, request_log_state: Optional[Dict[str, Any]] = None, extend_source_media_id: Optional[str] = None, ) -> AsyncGenerator: @@ -1826,7 +1837,6 @@ class GenerationHandler: import re as _re _uuid_match = _re.search(r'/video/([0-9a-f-]{36})', video_url or '') video_media_id = _uuid_match.group(1) if _uuid_match else video_info.get("mediaGenerationId", "") - print(f"[POLL-DEBUG] fifeUrl={str(video_url)[:80]}, video_media_id={video_media_id}") aspect_ratio = video_info.get("aspectRatio", "VIDEO_ASPECT_RATIO_LANDSCAPE") if not video_url: @@ -1862,7 +1872,14 @@ class GenerationHandler: # 递归轮询放大结果（不再放大） async for chunk in self._poll_video_result( - token, project_id, upsample_operations, stream, None, generation_result, response_state, request_log_state + token, + project_id, + upsample_operations, + stream, + None, + generation_result, + response_state, + request_log_state, ): yield chunk return @@ -1875,12 +1892,10 @@ class GenerationHandler: yield self._create_stream_chunk(f"⚠️ 放大失败: {str(e)}，返回原始视频\n") # ========== Extend 视频拼接 ========== - print(f"[CONCAT-DEBUG] extend_source_media_id={extend_source_media_id}, video_media_id={video_media_id}", flush=True) if extend_source_media_id and video_media_id: try: if stream: yield self._create_stream_chunk("\n视频续写完成，正在拼接完整视频...\n") - print(f"[CONCAT] 开始拼接: original={extend_source_media_id[:40]}..., extend={video_media_id[:40]}...", flush=True) debug_logger.log_info(f"[CONCAT] 开始拼接: original={extend_source_media_id[:12]}..., extend={video_media_id[:12]}...") # 提交拼接任务 @@ -1926,9 +1941,8 @@ class GenerationHandler: yield self._create_stream_chunk("⚠️ 拼接任务创建失败，返回续写片段\n") except Exception as e: import traceback - print(f"[CONCAT] ❌ 拼接异常: {str(e)}", flush=True) - print(f"[CONCAT] traceback: {traceback.format_exc()}", flush=True) debug_logger.log_error(f"[CONCAT] 拼接失败: {str(e)}") + debug_logger.log_error(f"[CONCAT] traceback: {traceback.format_exc()}") if stream: yield self._create_stream_chunk(f"⚠️ 拼接失败: {str(e)}，返回续写片段\n") # 拼接失败不影响返回，继续使用 extend 片段的 URL diff --git a/src/services/token_manager.py b/src/services/token_manager.py index 23b0e0f..133c0e3 100644 --- a/src/services/token_manager.py +++ b/src/services/token_manager.py @@ -209,7 +209,8 @@ class TokenManager: video_enabled: bool = True, image_concurrency: int = -1, video_concurrency: int = -1, - captcha_proxy_url: Optional[str] = None + captcha_proxy_url: Optional[str] = None, + extension_route_key: Optional[str] = None, ) -> Token: """Add a new token and prepare its pooled projects.""" existing_token = await self.db.get_token_by_st(st) @@ -284,7 +285,8 @@ class TokenManager: video_enabled=video_enabled, image_concurrency=image_concurrency, video_concurrency=video_concurrency, - captcha_proxy_url=captcha_proxy_url + captcha_proxy_url=captcha_proxy_url, + extension_route_key=extension_route_key, ) token_id = await self.db.add_token(token) @@ -314,7 +316,8 @@ class TokenManager: video_enabled: Optional[bool] = None, image_concurrency: Optional[int] = None, video_concurrency: Optional[int] = None, - captcha_proxy_url: Optional[str] = None + captcha_proxy_url: Optional[str] = None, + extension_route_key: Optional[str] = None, ): """Update token (支持修改project_id和project_name) @@ -344,6 +347,8 @@ class TokenManager: update_fields["video_concurrency"] = video_concurrency if captcha_proxy_url is not None: update_fields["captcha_proxy_url"] = captcha_proxy_url + if extension_route_key is not None: + update_fields["extension_route_key"] = extension_route_key # 检查token是否因429被禁用，如果是且未过期，则清空429状态 token = await self.db.get_token(token_id) diff --git a/test.py b/test.py deleted file mode 100644 index 3b21e89..0000000 --- a/test.py +++ /dev/null @@ -1,135 +0,0 @@ -import base64 -import json -import re -import urllib.request -import urllib.error -from pathlib import Path -from typing import Optional - -# ── 配置 ────────────────────────────────────────────── -IMAGE_PATH = Path(__file__).parent / "cat.jpg" -API_URL = "http://localhost:8000/v1/chat/completions" -API_KEY = "han1234" -MODEL = "gemini-3.1-flash-image-landscape" -PROMPT = "将这张图片变成水彩画风格" -OUTPUT = Path(__file__).parent / "cat_watercolor.jpg" -# ─────────────────────────────────────────────────────── - - -def read_image_base64(path: Path) -> str: - with open(path, "rb") as f: - return base64.b64encode(f.read()).decode() - - -def call_api(b64_image: str) -> str: - """发送请求，收集所有 SSE chunk，返回拼接后的完整 content。""" - payload = json.dumps({ - "model": MODEL, - "stream": True, - "messages": [ - { - "role": "user", - "content": [ - {"type": "text", "text": PROMPT}, - {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_image}"}} - ] - } - ] - }).encode() - - req = urllib.request.Request( - API_URL, - data=payload, - headers={ - "Authorization": f"Bearer {API_KEY}", - "Content-Type": "application/json", - }, - method="POST", - ) - - content = "" - reasoning = "" - with urllib.request.urlopen(req, timeout=300) as resp: - for raw_line in resp: - line = raw_line.decode("utf-8", errors="replace").strip() - if not line.startswith("data:"): - continue - data_str = line[5:].strip() - if data_str == "[DONE]": - break - try: - chunk = json.loads(data_str) - except json.JSONDecodeError: - continue - - # 检查错误 - if "error" in chunk: - raise RuntimeError(chunk["error"].get("message", str(chunk["error"]))) - - delta = chunk.get("choices", [{}])[0].get("delta", {}) - if delta.get("reasoning_content"): - r = delta["reasoning_content"] - reasoning += r - print(r, end="", flush=True) - if delta.get("content"): - content += delta["content"] - - print() # 换行 - return content - - -def extract_image_url(content: str) -> Optional[str]: - """从 markdown 或纯文本中提取图片 URL。""" - # ![...](url) 或直接 URL - patterns = [ - r"!\[.*?\]\((https?://\S+?)\)", - r"(https?://\S+\.(?:jpg|jpeg|png|webp|gif)(?:\?\S*)?)", - r"(https?://\S+)", - ] - for pat in patterns: - m = re.search(pat, content) - if m: - return m.group(1) - return None - - -def download_image(url: str, dest: Path): - print(f"📥 下载图片: {url}") - urllib.request.urlretrieve(url, dest) - print(f"✅ 已保存到: {dest}") - - -def main(): - if not IMAGE_PATH.exists(): - raise FileNotFoundError(f"找不到输入图片: {IMAGE_PATH}") - - print(f"📷 读取图片: {IMAGE_PATH}") - b64 = read_image_base64(IMAGE_PATH) - - print(f"🚀 发送请求到 {API_URL} ...") - print("-" * 50) - content = call_api(b64) - print("-" * 50) - - if not content: - print("⚠️ 响应 content 为空，可能图片已在 reasoning_content 中描述，请查看上方日志。") - return - - print(f"\n📝 响应内容:\n{content}\n") - - url = extract_image_url(content) - if url: - download_image(url, OUTPUT) - else: - # 尝试 base64 内嵌 - b64_match = re.search(r"data:image/[^;]+;base64,([A-Za-z0-9+/=]+)", content) - if b64_match: - img_data = base64.b64decode(b64_match.group(1)) - OUTPUT.write_bytes(img_data) - print(f"✅ base64 图片已保存到: {OUTPUT}") - else: - print("⚠️ 未能从响应中提取图片 URL，请手动查看上方 content。") - - -if __name__ == "__main__": - main() diff --git a/tests/test_veo_lite_support.py b/tests/test_veo_lite_support.py index 4779f5f..2e57662 100644 --- a/tests/test_veo_lite_support.py +++ b/tests/test_veo_lite_support.py @@ -62,7 +62,7 @@ class VeoLiteFlowClientTests(unittest.IsolatedAsyncioTestCase): async def test_generate_video_text_uses_v2_payload_for_lite(self): captured = {} - async def fake_make_request(method, url, json_data, use_at, at_token): + async def fake_make_request(method, url, json_data, use_at, at_token, **kwargs): captured["url"] = url captured["json_data"] = json_data return {"operations": [{"operation": {"name": "task-1"}}]} @@ -92,7 +92,7 @@ class VeoLiteFlowClientTests(unittest.IsolatedAsyncioTestCase): async def test_generate_video_start_end_uses_v2_payload_for_interpolation_lite(self): captured = {} - async def fake_make_request(method, url, json_data, use_at, at_token): + async def fake_make_request(method, url, json_data, use_at, at_token, **kwargs): captured["url"] = url captured["json_data"] = json_data return {"operations": [{"operation": {"name": "task-2"}}]}