diff --git a/services/kdocs_uploader.py b/services/kdocs_uploader.py index 4d0ef41..6c87a02 100644 --- a/services/kdocs_uploader.py +++ b/services/kdocs_uploader.py @@ -199,11 +199,31 @@ class KDocsUploader: pass self._playwright = None - def _open_document(self, doc_url: str) -> bool: + def _open_document(self, doc_url: str, *, fast: bool = False) -> bool: try: self._doc_url = doc_url - self._page.goto(doc_url) - time.sleep(1) + if fast: + doc_pages = self._find_doc_pages(doc_url) + if doc_pages: + self._page = doc_pages[0] + return True + login_pages = [] + for page in self._list_pages(): + try: + url = getattr(page, "url", "") or "" + if self._is_login_url(url): + login_pages.append(page) + except Exception: + continue + if login_pages: + self._page = login_pages[0] + return True + goto_kwargs = {} + if fast: + fast_timeout = int(os.environ.get("KDOCS_FAST_GOTO_TIMEOUT_MS", "15000")) + goto_kwargs = {"wait_until": "domcontentloaded", "timeout": fast_timeout} + self._page.goto(doc_url, **goto_kwargs) + time.sleep(0.6) doc_pages = self._find_doc_pages(doc_url) if doc_pages and doc_pages[0] is not self._page: self._page = doc_pages[0] @@ -315,22 +335,40 @@ class KDocsUploader: storage_state = getattr(config, "KDOCS_LOGIN_STATE_FILE", "data/kdocs_login_state.json") return os.path.exists(storage_state) - def _ensure_login_dialog(self) -> None: + def _ensure_login_dialog( + self, + *, + timeout_ms: int = 1200, + frame_timeout_ms: int = 800, + quick: bool = False, + ) -> None: login_names = ["登录并加入编辑", "立即登录", "登录"] wechat_names = ["微信登录", "微信扫码登录", "微信扫码", "扫码登录"] pages = self._iter_pages() clicked = False for page in pages: - if self._try_click_names(page, login_names): + if self._try_click_names( + page, + login_names, + timeout_ms=timeout_ms, + frame_timeout_ms=frame_timeout_ms, + quick=quick, + ): clicked = True break if clicked: time.sleep(1.5) pages = self._iter_pages() for page in pages: - if self._try_click_names(page, wechat_names): + if self._try_click_names( + page, + wechat_names, + timeout_ms=timeout_ms, + frame_timeout_ms=frame_timeout_ms, + quick=quick, + ): return - self._try_confirm_login() + self._try_confirm_login(timeout_ms=timeout_ms, frame_timeout_ms=frame_timeout_ms, quick=quick) def _capture_qr_image(self) -> Optional[bytes]: pages = self._iter_pages() @@ -458,34 +496,44 @@ class KDocsUploader: return False return False - def _try_click_names(self, page, names: list) -> bool: + def _try_click_names( + self, + page, + names: list, + *, + timeout_ms: int = 1200, + frame_timeout_ms: int = 800, + quick: bool = False, + ) -> bool: for name in names: - if self._try_click_role(page, "button", name, timeout=1200): - return True - if self._try_click_role(page, "link", name, timeout=1200): + if self._try_click_role(page, "button", name, timeout=timeout_ms): return True + if not quick: + if self._try_click_role(page, "link", name, timeout=timeout_ms): + return True try: el = page.get_by_text(name, exact=True) - if el.is_visible(timeout=1200): - el.click() - time.sleep(1) - return True - except Exception: - pass - try: - el = page.get_by_text(name, exact=False) - if el.is_visible(timeout=1200): + if el.is_visible(timeout=timeout_ms): el.click() time.sleep(1) return True except Exception: pass + if not quick: + try: + el = page.get_by_text(name, exact=False) + if el.is_visible(timeout=timeout_ms): + el.click() + time.sleep(1) + return True + except Exception: + pass try: for frame in page.frames: for name in names: try: el = frame.get_by_role("button", name=name) - if el.is_visible(timeout=800): + if el.is_visible(timeout=frame_timeout_ms): el.click() time.sleep(1) return True @@ -493,29 +541,42 @@ class KDocsUploader: pass try: el = frame.get_by_text(name, exact=True) - if el.is_visible(timeout=800): - el.click() - time.sleep(1) - return True - except Exception: - pass - try: - el = frame.get_by_text(name, exact=False) - if el.is_visible(timeout=800): + if el.is_visible(timeout=frame_timeout_ms): el.click() time.sleep(1) return True except Exception: pass + if not quick: + try: + el = frame.get_by_text(name, exact=False) + if el.is_visible(timeout=frame_timeout_ms): + el.click() + time.sleep(1) + return True + except Exception: + pass except Exception: return False return False - def _try_confirm_login(self) -> bool: + def _try_confirm_login( + self, + *, + timeout_ms: int = 1200, + frame_timeout_ms: int = 800, + quick: bool = False, + ) -> bool: confirm_names = ["确认登录", "确认登陆"] pages = self._iter_pages() for page in pages: - if self._try_click_names(page, confirm_names): + if self._try_click_names( + page, + confirm_names, + timeout_ms=timeout_ms, + frame_timeout_ms=frame_timeout_ms, + quick=quick, + ): return True return False @@ -602,7 +663,7 @@ class KDocsUploader: self._handle_clear_login() if not self._ensure_playwright(use_storage_state=not force): return {"success": False, "error": self._last_error or "浏览器不可用"} - if not self._open_document(doc_url): + if not self._open_document(doc_url, fast=True): return {"success": False, "error": self._last_error or "打开文档失败"} if not force and self._has_saved_login_state() and self._is_logged_in(): @@ -611,11 +672,21 @@ class KDocsUploader: self._save_login_state() return {"success": True, "logged_in": True, "qr_image": ""} - self._ensure_login_dialog() + fast_login_timeout = int(os.environ.get("KDOCS_FAST_LOGIN_TIMEOUT_MS", "300")) + self._ensure_login_dialog( + timeout_ms=fast_login_timeout, + frame_timeout_ms=fast_login_timeout, + quick=True, + ) qr_image = None invalid_qr = None - for _ in range(10): - self._ensure_login_dialog() + for attempt in range(10): + if attempt in (3, 7): + self._ensure_login_dialog( + timeout_ms=fast_login_timeout, + frame_timeout_ms=fast_login_timeout, + quick=True, + ) candidate = self._capture_qr_image() if candidate and self._is_valid_qr_image(candidate): qr_image = candidate @@ -689,10 +760,19 @@ class KDocsUploader: return {"success": True, "logged_in": False, "error": "未配置文档链接"} if not self._ensure_playwright(): return {"success": False, "logged_in": False, "error": self._last_error or "浏览器不可用"} - if not self._open_document(doc_url): + if not self._open_document(doc_url, fast=True): return {"success": False, "logged_in": False, "error": self._last_error or "打开文档失败"} - self._ensure_login_dialog() - self._try_confirm_login() + fast_login_timeout = int(os.environ.get("KDOCS_FAST_LOGIN_TIMEOUT_MS", "300")) + self._ensure_login_dialog( + timeout_ms=fast_login_timeout, + frame_timeout_ms=fast_login_timeout, + quick=True, + ) + self._try_confirm_login( + timeout_ms=fast_login_timeout, + frame_timeout_ms=fast_login_timeout, + quick=True, + ) logged_in = self._is_logged_in() self._last_login_ok = logged_in self._login_required = not logged_in