diff --git a/Dockerfile b/Dockerfile index 95a7a0b..dd5ebe2 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,10 +5,19 @@ WORKDIR /app COPY pyproject.toml . COPY src ./src -RUN pip install --no-cache-dir ".[tor,instsci]" +RUN pip install --no-cache-dir ".[tor,cloakbrowser,instsci]" FROM python:3.12-slim +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + libglib2.0-0 libnss3 libnspr4 libdbus-1-3 \ + libatk1.0-0 libatk-bridge2.0-0 libcups2 libdrm2 \ + libxkbcommon0 libxcomposite1 libxdamage1 libxfixes3 \ + libxrandr2 libgbm1 libpango-1.0-0 libcairo2 libasound2 \ + libatspi2.0-0 libwayland-client0 tini \ + && rm -rf /var/lib/apt/lists/* + WORKDIR /app COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages @@ -19,6 +28,9 @@ COPY pyproject.toml . EXPOSE 8000 ENV SCANSCI_PDF_DATA_DIR=/data/paper-fetch +ENV CLOAKBROWSER_CACHE_DIR=/data/paper-fetch/browser-cache ENV MCP_MODE=streamable_http +ENV MALLOC_ARENA_MAX=2 +ENTRYPOINT ["tini", "--"] CMD ["python", "-m", "scansci_pdf", "run", "--mode", "streamable_http"] diff --git a/pyproject.toml b/pyproject.toml index fa734d1..f0c817f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "requests>=2.31", "requests[socks]>=2.31", "beautifulsoup4>=4.12", + "PyMuPDF>=1.24", "mcp[cli]>=1.12", "typer>=0.15", "uvicorn>=0.34", diff --git a/src/scansci_pdf/auth.py b/src/scansci_pdf/auth.py index fa9331d..90d36c8 100644 --- a/src/scansci_pdf/auth.py +++ b/src/scansci_pdf/auth.py @@ -10,14 +10,14 @@ import requests from Crypto.Cipher import AES -from .cloakbrowser_compat import prepare_cloakbrowser_runtime +from .cloakbrowser_compat import launch_with_driver_cleanup, prepare_cloakbrowser_runtime from .session_store import CookieStore try: prepare_cloakbrowser_runtime() from cloakbrowser import launch _HAS_CLOAKBROWSER = True -except ImportError: +except Exception: launch = None # type: ignore[assignment] _HAS_CLOAKBROWSER = False @@ -65,6 +65,7 @@ def __init__( self._browser = None self._context = None self._page = None + self._browser_lease = None base = config.get("instsci_base_url", "") self._webvpn_base = base.rstrip("/") if base else "" @@ -158,6 +159,7 @@ def _validate_session(self) -> bool: test_url = TEST_URL else: test_url = self.convert_url(TEST_URL) + resp = None try: resp = self.session.get(test_url, timeout=15, allow_redirects=True) if "cas" in resp.url.lower() or "login" in resp.url.lower(): @@ -167,6 +169,12 @@ def _validate_session(self) -> bool: return True except requests.RequestException as e: logger.warning("Session validation failed: %s", e) + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass return False def _browser_login(self) -> bool: @@ -174,24 +182,41 @@ def _browser_login(self) -> bool: logger.error("cloakbrowser not installed. Run: pip install cloakbrowser") return False + if not self._close_browser(): + logger.error("Previous WebVPN browser could not be closed; refusing replacement") + return False try: + from . import browser_engine + self._browser_lease = browser_engine._retain_browser_slot(self.config) prepare_cloakbrowser_runtime() from cloakbrowser import launch_persistent_context profile_dir = _get_profile_dir(self.config) profile_dir.mkdir(parents=True, exist_ok=True) - self._context = launch_persistent_context( + raw_context = launch_with_driver_cleanup( + launch_persistent_context, user_data_dir=str(profile_dir), headless=False, humanize=True, args=self._browser_launch_args(), ) + self._context = browser_engine._LeasedPersistentContext( + raw_context, + self._browser_lease, + ) + self._browser_lease = None self._browser = None self._page = self._context.new_page() except Exception as e: logger.error("Failed to start CloakBrowser: %s", e) + self._close_browser() return False - self._page.goto(self._webvpn_base, wait_until="networkidle", timeout=30000) - current_url = self._page.url + try: + self._page.goto(self._webvpn_base, wait_until="networkidle", timeout=30000) + current_url = self._page.url + except Exception as e: + logger.error("Failed to open campus gateway: %s", e) + self._close_browser() + return False logger.info("Session test: navigated to campus gateway, landed on %s", current_url[:80]) parsed = urlparse(current_url) @@ -208,10 +233,20 @@ def _browser_login(self) -> bool: if not on_login_page and not is_idp: logger.info("Persistent context has valid session! URL=%s", current_url[:60]) - self._save_browser_cookies() + try: + self._save_browser_cookies() + except Exception as e: + logger.error("Failed to save browser cookies: %s", e) + self._close_browser() + return False return True - self._page.goto(self._webvpn_base, wait_until="domcontentloaded") + try: + self._page.goto(self._webvpn_base, wait_until="domcontentloaded") + except Exception as e: + logger.error("Failed to open campus login page: %s", e) + self._close_browser() + return False print("\n" + "=" * 60) print(f" Please log in at {self._webvpn_base}") @@ -231,9 +266,7 @@ def _browser_login(self) -> bool: try: if not self._context.pages: logger.info("Browser closed by user.") - self._browser = None - self._context = None - self._page = None + self._close_browser() return False current_url = self._page.url @@ -280,9 +313,7 @@ def _browser_login(self) -> bool: except Exception: logger.warning("Browser connection lost.") - self._browser = None - self._context = None - self._page = None + self._close_browser() return False print("\n Login timed out after 10 minutes.\n") @@ -298,20 +329,46 @@ def _save_browser_cookies(self): logger.info("Saved %d cookies to %s", len(cookies), store.path) store.apply_to_session(self.session, cookies) - def _close_browser(self): - if self._context: - try: - self._context.close() - except Exception: - pass - if self._browser: - try: - self._browser.close() - except Exception: - pass - self._browser = None - self._context = None - self._page = None + def _close_browser(self) -> bool: + from . import browser_engine + + context = self._context + browser = self._browser + lease = self._browser_lease + if lease is not None: + lease._assert_owner() + + if context is None: + context_closed, context_errors = True, [] + else: + context_closed, context_errors = browser_engine._close_resource_with_confirmation( + context + ) + if browser is None: + browser_closed, browser_errors = True, [] + else: + browser_closed, browser_errors = browser_engine._close_resource_with_confirmation( + browser, + is_browser=True, + ) + if browser_closed: + context_closed = True + + shutdown_complete = context_closed and browser_closed + if shutdown_complete: + self._browser = None + self._context = None + self._page = None + self._browser_lease = None + if lease is not None: + lease.close() + else: + errors = [str(error) for error in context_errors + browser_errors] + logger.warning( + "WebVPN browser close incomplete; retaining resources and global slot: %s", + "; ".join(errors), + ) + return shutdown_complete def fetch(self, url: str, **kwargs) -> requests.Response: """Fetch a URL through the campus, EasyConnect, or connector session.""" @@ -331,9 +388,10 @@ def fetch(self, url: str, **kwargs) -> requests.Response: def close(self): self._close_browser() - if self._session: - self._session.close() - self._session = None + session = self._session + self._session = None + if session is not None: + session.close() class EZProxyAuth: @@ -350,11 +408,18 @@ def __init__( self._browser = None self._context = None self._page = None + self._browser_lease = None @property def browser_context(self): return self._context + def _browser_launch_args(self) -> list[str]: + return [ + "--no-proxy-server", + "--disable-features=CrossOriginOpenerPolicy", + ] + @property def session(self) -> requests.Session: if self._session is None: @@ -382,6 +447,7 @@ def _try_load_cookies(self) -> bool: return self._validate_session() def _validate_session(self) -> bool: + resp = None try: resp = self.session.get(self._proxy_base + TEST_URL, timeout=15, allow_redirects=True) if "login" in resp.url.lower() or "cas" in resp.url.lower(): @@ -389,24 +455,47 @@ def _validate_session(self) -> bool: return resp.status_code == 200 except requests.RequestException: return False + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass def _browser_login(self) -> bool: if not _HAS_CLOAKBROWSER: logger.error("cloakbrowser not installed. Run: pip install cloakbrowser") return False + if not self._close_browser(): + logger.error("Previous EZproxy browser could not be closed; refusing replacement") + return False try: - self._browser = launch( + from . import browser_engine + self._browser_lease = browser_engine._retain_browser_slot(self.config) + raw_browser = launch_with_driver_cleanup( + launch, headless=False, humanize=True, args=self._browser_launch_args(), ) + self._browser = browser_engine._LeasedBrowser( + raw_browser, + self._browser_lease, + ) + self._browser_lease = None self._context = self._browser.new_context() self._page = self._context.new_page() except Exception as e: logger.error("Failed to start CloakBrowser: %s", e) + self._close_browser() return False - self._page.goto(self._proxy_base + TEST_URL, wait_until="domcontentloaded") + try: + self._page.goto(self._proxy_base + TEST_URL, wait_until="domcontentloaded") + except Exception as e: + logger.error("Failed to open EZproxy login page: %s", e) + self._close_browser() + return False print("\n" + "=" * 60) print(f" Please log in at the EZproxy page.") @@ -425,9 +514,7 @@ def _browser_login(self) -> bool: try: if not self._context.pages: logger.info("Browser closed by user.") - self._browser = None - self._context = None - self._page = None + self._close_browser() return False current_url = self._page.url @@ -445,9 +532,7 @@ def _browser_login(self) -> bool: except Exception: logger.warning("Browser connection lost.") - self._browser = None - self._context = None - self._page = None + self._close_browser() return False print("\n Login timed out after 10 minutes.\n") @@ -463,15 +548,46 @@ def _save_browser_cookies(self): logger.info("Saved %d cookies to %s", len(cookies), store.path) store.apply_to_session(self.session, cookies) - def _close_browser(self): - if self._browser: - try: - self._browser.close() - except Exception: - pass + def _close_browser(self) -> bool: + from . import browser_engine + + context = self._context + browser = self._browser + lease = self._browser_lease + if lease is not None: + lease._assert_owner() + + if context is None: + context_closed, context_errors = True, [] + else: + context_closed, context_errors = browser_engine._close_resource_with_confirmation( + context + ) + if browser is None: + browser_closed, browser_errors = True, [] + else: + browser_closed, browser_errors = browser_engine._close_resource_with_confirmation( + browser, + is_browser=True, + ) + if browser_closed: + context_closed = True + + shutdown_complete = context_closed and browser_closed + if shutdown_complete: self._browser = None self._context = None self._page = None + self._browser_lease = None + if lease is not None: + lease.close() + else: + errors = [str(error) for error in context_errors + browser_errors] + logger.warning( + "EZproxy browser close incomplete; retaining resources and global slot: %s", + "; ".join(errors), + ) + return shutdown_complete def get_proxied_url(self, url: str) -> str: if self._proxy_base and self._proxy_base.rstrip("/").split("//")[-1].split("/")[0] in url: @@ -486,6 +602,7 @@ def fetch(self, url: str, **kwargs) -> requests.Response: def close(self): self._close_browser() - if self._session: - self._session.close() - self._session = None + session = self._session + self._session = None + if session is not None: + session.close() diff --git a/src/scansci_pdf/browser_cookies.py b/src/scansci_pdf/browser_cookies.py index bf7bcf7..52e9fa3 100644 --- a/src/scansci_pdf/browser_cookies.py +++ b/src/scansci_pdf/browser_cookies.py @@ -109,8 +109,14 @@ def extract_via_browser( Result dict with success, cookies_count, domains, etc. """ try: + from .cloakbrowser_compat import ( + launch_with_driver_cleanup, + prepare_cloakbrowser_runtime, + ) + + prepare_cloakbrowser_runtime() from cloakbrowser import launch - except ImportError: + except Exception: return { "success": False, "error": "cloakbrowser not installed", @@ -127,8 +133,20 @@ def extract_via_browser( print(f" 打开页面: {url}") print(f" 登录完成后关闭浏览器窗口即可\n") + browser = None + context = None + page = None + slot_lease = None try: - browser = launch(headless=False, humanize=True) + from . import browser_engine + slot_lease = browser_engine._retain_browser_slot(config) + raw_browser = launch_with_driver_cleanup( + launch, + headless=False, + humanize=True, + ) + browser = browser_engine._LeasedBrowser(raw_browser, slot_lease) + slot_lease = None context = browser.new_context(viewport={"width": 1440, "height": 900}) page = context.new_page() @@ -163,7 +181,6 @@ def extract_via_browser( all_cookies = context.cookies() if not all_cookies: - browser.close() return { "success": False, "message": "未捕获到 cookies。请确保已登录机构账号。", @@ -186,8 +203,6 @@ def extract_via_browser( except Exception: pass - browser.close() - domains_found = list({c.get("domain", "").lstrip(".") for c in save_cookies})[:10] return { "success": True, @@ -204,6 +219,31 @@ def extract_via_browser( except Exception as exc: log.info(f" [cookies] Error: {exc}") return {"success": False, "error": str(exc)} + finally: + from . import browser_engine + if page is not None: + try: + page.close() + except Exception: + pass + context_closed = context is None + if context is not None: + context_closed, _ = browser_engine._close_resource_with_confirmation(context) + browser_closed = browser is None + if browser is not None: + browser_closed, close_errors = browser_engine._close_resource_with_confirmation( + browser, + is_browser=True, + ) + if browser_closed: + context_closed = True + elif close_errors: + log.info( + " [cookies] Browser close failed; retaining global slot: " + + "; ".join(str(error) for error in close_errors) + ) + if slot_lease is not None and context_closed and browser_closed: + slot_lease.close() def load_saved_cookies(config: dict[str, Any]) -> list[dict[str, Any]]: diff --git a/src/scansci_pdf/browser_engine.py b/src/scansci_pdf/browser_engine.py index d2adfc2..48de4af 100644 --- a/src/scansci_pdf/browser_engine.py +++ b/src/scansci_pdf/browser_engine.py @@ -13,6 +13,7 @@ from __future__ import annotations import base64 +import contextlib import json import logging import time @@ -30,8 +31,22 @@ _HAS_CLOAKBROWSER: bool | None = None +def _prepare_cloakbrowser_runtime() -> bool: + """Prepare CloakBrowser before importing any of its launch APIs.""" + try: + from .cloakbrowser_compat import prepare_cloakbrowser_runtime + + prepare_cloakbrowser_runtime() + except Exception as exc: + logger.warning("browser_engine: CloakBrowser runtime preparation failed: %s", exc) + return False + return True + + def _check_cloakbrowser() -> bool: global _HAS_CLOAKBROWSER + if not _prepare_cloakbrowser_runtime(): + return False if _HAS_CLOAKBROWSER is None: try: from cloakbrowser import launch # noqa: F401 @@ -51,6 +66,497 @@ def _check_cloakbrowser() -> bool: import threading as _threading _tls = _threading.local() +_browser_semaphore: _threading.Semaphore | None = None +_browser_semaphore_lock = _threading.Lock() +_browser_semaphore_capacity = 0 +_browser_active = 0 +_browser_waiters = 0 +_browser_shutdown_generation = 0 +_browser_reclaimed_generation = 0 +_browser_memory_reclaim_lock = _threading.Lock() +_retained_browser_resources: dict[_threading.Thread, list[Any]] = {} +_retained_browser_resources_lock = _threading.RLock() +MAX_BROWSER_WORKERS = 4 + + +class BrowserOperationCancelled(RuntimeError): + """Raised when queued browser work is cancelled before launch.""" + + +def _register_retained_browser_resource(resource: Any) -> None: + """Keep an unclosed Playwright handle reachable for its owner-thread retry.""" + owner = _threading.current_thread() + with _retained_browser_resources_lock: + resources = _retained_browser_resources.setdefault(owner, []) + if not any(candidate is resource for candidate in resources): + resources.append(resource) + + +def _forget_retained_browser_resource(resource: Any) -> None: + """Forget a retained handle only after its close has been confirmed.""" + owner = getattr(resource, "_owner_thread", _threading.current_thread()) + with _retained_browser_resources_lock: + resources = _retained_browser_resources.get(owner) + if not resources: + return + resources[:] = [candidate for candidate in resources if candidate is not resource] + if not resources: + _retained_browser_resources.pop(owner, None) + + +def _owner_retained_browser_resources() -> tuple[Any, ...]: + # Key by the Thread object, not its reusable numeric ident. If an owner + # exits while a driver remains live, a later thread must not touch that + # Playwright handle merely because the OS recycled the same identifier. + owner = _threading.current_thread() + with _retained_browser_resources_lock: + return tuple(_retained_browser_resources.get(owner, ())) + + +def _retry_retained_browser_resources() -> tuple[bool, list[Exception]]: + """Retry failed closes without ever crossing Playwright thread ownership.""" + errors: list[Exception] = [] + for resource in _owner_retained_browser_resources(): + try: + resource.close() + except Exception as exc: + # A driver may raise after it has nevertheless confirmed closure. + # In that case the proxy has already removed itself from the registry. + if any( + candidate is resource + for candidate in _owner_retained_browser_resources() + ): + errors.append(exc) + return not _owner_retained_browser_resources(), errors + + +class _BrowserSlotToken: + """One process-wide browser permit shared by nested work on its owner thread.""" + + def __init__(self, semaphore: _threading.Semaphore): + self.semaphore = semaphore + self.owner_thread = _threading.current_thread() + self.owner_thread_id = _threading.get_ident() + self.references = 1 + self.released = False + + +class BrowserSlotLease: + """Owner-thread lease for one slot in the process-wide browser limiter. + + Scoped leases are re-entrant on one thread: a source worker may hold the + outer lease while ``_get_shared_browser()`` or ``get_persistent_context()`` + retain the same permit for the actual Playwright resource. Resource leases + are not left on the re-entrancy stack after their creator returns, so a + leaked/failed-to-close browser cannot let later work bypass the limiter. + """ + + def __init__( + self, + token: _BrowserSlotToken, + *, + scoped: bool, + ) -> None: + self._token = token + self._owner_thread = token.owner_thread + self._owner_thread_id = token.owner_thread_id + self._scoped = scoped + self._closed = False + + @property + def semaphore(self) -> _threading.Semaphore: + return self._token.semaphore + + def _assert_owner(self) -> None: + current_thread = _threading.current_thread() + current = _threading.get_ident() + if current_thread is not self._owner_thread: + raise RuntimeError( + "browser slot lease belongs to another thread; " + f"owner={self._owner_thread_id}, current={current}" + ) + + def close(self) -> None: + """Release this reference, and the permit only after the last reference.""" + self._assert_owner() + if self._closed: + return + + if self._scoped: + stack = getattr(_tls, "browser_slot_stack", None) + if not stack or stack[-1] is not self: + raise RuntimeError("browser slot scopes must close in owner-thread LIFO order") + stack.pop() + + release_semaphore = None + with _browser_semaphore_lock: + if self._token.released or self._token.references <= 0: + raise RuntimeError("browser slot token was already released") + self._token.references -= 1 + if self._token.references == 0: + self._token.released = True + release_semaphore = self._token.semaphore + self._closed = True + if release_semaphore is not None: + _release_browser_slot(release_semaphore) + + def __enter__(self): + self._assert_owner() + if self._closed: + raise RuntimeError("browser slot lease is already closed") + return self + + def __exit__(self, exc_type, exc, traceback) -> None: + self.close() + + +def _new_browser_slot_lease( + config: dict[str, Any] | None, + cancel_event: _threading.Event | None, + *, + scoped: bool, +) -> BrowserSlotLease: + """Acquire or re-enter the global limiter on the current owner thread.""" + cancel_event = _effective_cancel_event(cancel_event) + if _cancelled(cancel_event): + raise BrowserOperationCancelled("browser operation cancelled before launch") + + # A previous direct/persistent resource can have failed to close after its + # caller dropped the local variable. Retry that exact owner-thread handle + # before taking/re-entering a permit; otherwise max=1 would deadlock and a + # re-entrant outer source lease could launch a replacement beside it. + if _owner_retained_browser_resources(): + retained_closed, retained_errors = _retry_retained_browser_resources() + if not retained_closed: + detail = "; ".join(str(error) for error in retained_errors) + raise RuntimeError( + "previous owner-thread browser resource could not be closed; " + "replacement launch refused" + + (f": {detail}" if detail else "") + ) + + stack = getattr(_tls, "browser_slot_stack", None) + if stack is None: + stack = [] + _tls.browser_slot_stack = stack + + if stack: + token = stack[-1]._token + if token.owner_thread is not _threading.current_thread(): + raise RuntimeError("browser slot scope crossed thread ownership") + with _browser_semaphore_lock: + if token.released: + raise RuntimeError("cannot re-enter a released browser slot") + token.references += 1 + else: + semaphore = _prepare_browser_slot(config) + _acquire_browser_slot( + semaphore, + waiter_registered=True, + cancel_event=cancel_event, + ) + token = _BrowserSlotToken(semaphore) + + lease = BrowserSlotLease(token, scoped=scoped) + if scoped: + stack.append(lease) + return lease + + +def browser_slot( + config: dict[str, Any] | None = None, + cancel_event: _threading.Event | None = None, +) -> BrowserSlotLease: + """Return a re-entrant scoped lease for a browser-backed operation.""" + return _new_browser_slot_lease(config, cancel_event, scoped=True) + + +def _retain_browser_slot( + config: dict[str, Any] | None = None, + cancel_event: _threading.Event | None = None, +) -> BrowserSlotLease: + """Retain a permit until an owning browser/context is confirmed closed.""" + return _new_browser_slot_lease(config, cancel_event, scoped=False) + + +class _LeasedPersistentContext: + """Proxy that releases a browser slot when its context is closed.""" + + def __init__( + self, + context: Any, + lease: BrowserSlotLease | _threading.Semaphore, + ): + self._context = context + # Accept a raw semaphore for compatibility with low-level lifecycle + # tests; production paths always pass an owner-thread lease. + self._lease = lease if isinstance(lease, BrowserSlotLease) else None + self._semaphore = lease.semaphore if self._lease is not None else lease + self._owner_thread = _threading.current_thread() + self._owner_thread_id = _threading.get_ident() + self._closed = False + self._close_lock = _threading.Lock() + + def __getattr__(self, name: str) -> Any: + if self._lease is not None: + self._lease._assert_owner() + return getattr(self._context, name) + + def close(self) -> None: + with self._close_lock: + if self._closed: + return + if self._lease is not None: + self._lease._assert_owner() + try: + self._context.close() + except Exception: + if not _context_is_confirmed_closed(self._context): + _register_retained_browser_resource(self) + raise + self._closed = True + self._release_lease() + _forget_retained_browser_resource(self) + raise + self._closed = True + self._release_lease() + _forget_retained_browser_resource(self) + + def _release_lease(self) -> None: + if self._lease is not None: + self._lease.close() + else: + _release_browser_slot(self._semaphore) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback) -> None: + self.close() + + +class _LeasedBrowser: + """Browser proxy that owns a slot until the process is confirmed closed.""" + + def __init__(self, browser: Any, lease: BrowserSlotLease): + self._browser = browser + self._lease = lease + self._owner_thread = _threading.current_thread() + self._owner_thread_id = _threading.get_ident() + self._closed = False + self._close_lock = _threading.Lock() + + def __getattr__(self, name: str) -> Any: + self._lease._assert_owner() + return getattr(self._browser, name) + + def __eq__(self, other: Any) -> bool: + if isinstance(other, _LeasedBrowser): + other = other._browser + return self._browser == other + + def close(self) -> None: + with self._close_lock: + if self._closed: + return + self._lease._assert_owner() + closed, errors = _close_resource_with_confirmation( + self._browser, + is_browser=True, + ) + if not closed: + _register_retained_browser_resource(self) + if errors: + raise errors[-1] + raise RuntimeError("browser close could not be confirmed") + self._closed = True + self._lease.close() + _forget_retained_browser_resource(self) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback) -> None: + self.close() + + +def _browser_worker_limit(config: dict[str, Any] | None) -> int: + try: + configured = int((config or {}).get("max_browser_workers", 1)) + except (TypeError, ValueError): + return 1 + return min(MAX_BROWSER_WORKERS, max(1, configured)) + + +def _get_browser_semaphore(config: dict[str, Any] | None) -> _threading.Semaphore: + desired = _browser_worker_limit(config) + with _browser_semaphore_lock: + return _select_browser_semaphore_locked(desired) + + +def _select_browser_semaphore_locked(desired: int) -> _threading.Semaphore: + """Select the limiter while ``_browser_semaphore_lock`` is held.""" + global _browser_semaphore, _browser_semaphore_capacity + if _browser_semaphore is None or ( + _browser_active == 0 + and _browser_waiters == 0 + and _browser_semaphore_capacity != desired + ): + _browser_semaphore = _threading.Semaphore(desired) + _browser_semaphore_capacity = desired + return _browser_semaphore + + +def _prepare_browser_slot(config: dict[str, Any] | None) -> _threading.Semaphore: + """Select a limiter and register its waiter as one atomic operation.""" + global _browser_waiters + desired = _browser_worker_limit(config) + with _browser_semaphore_lock: + semaphore = _select_browser_semaphore_locked(desired) + _browser_waiters += 1 + return semaphore + + +def _set_thread_cancel_event(cancel_event: _threading.Event | None): + """Set cooperative cancellation for browser work on the current thread.""" + previous = getattr(_tls, "cancel_event", None) + _tls.cancel_event = cancel_event + return previous + + +def _effective_cancel_event( + cancel_event: _threading.Event | None = None, +) -> _threading.Event | None: + return cancel_event if cancel_event is not None else getattr(_tls, "cancel_event", None) + + +def _cancelled(cancel_event: _threading.Event | None) -> bool: + return cancel_event is not None and cancel_event.is_set() + + +def _wait_or_cancel( + cancel_event: _threading.Event | None, + timeout: float, +) -> bool: + if cancel_event is None: + time.sleep(max(0.0, timeout)) + return False + return cancel_event.wait(max(0.0, timeout)) + + +def _write_pdf_bytes_atomic( + output_path: Path, + pdf_bytes: bytes, + cancel_event: _threading.Event | None = None, +) -> bool: + """Replace a PDF only after a complete, non-cancelled temporary write.""" + from .pdf_utils import write_pdf_bytes_atomic + + return write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event) + + +def _acquire_browser_slot( + semaphore: _threading.Semaphore, + *, + waiter_registered: bool = False, + cancel_event: _threading.Event | None = None, +) -> None: + global _browser_active, _browser_waiters + if not waiter_registered: + with _browser_semaphore_lock: + _browser_waiters += 1 + try: + while True: + effective_cancel = _effective_cancel_event(cancel_event) + if effective_cancel is not None and effective_cancel.is_set(): + raise BrowserOperationCancelled("browser operation cancelled before launch") + if not semaphore.acquire(timeout=0.1): + continue + effective_cancel = _effective_cancel_event(cancel_event) + if effective_cancel is not None and effective_cancel.is_set(): + semaphore.release() + raise BrowserOperationCancelled("browser operation cancelled before launch") + with _browser_semaphore_lock: + _browser_active += 1 + return + finally: + with _browser_semaphore_lock: + _browser_waiters = max(0, _browser_waiters - 1) + + +def _release_browser_slot(semaphore: _threading.Semaphore) -> None: + global _browser_active, _browser_shutdown_generation + semaphore.release() + with _browser_semaphore_lock: + _browser_active = max(0, _browser_active - 1) + if _browser_active == 0: + _browser_shutdown_generation += 1 + + +def _trim_process_heap() -> bool: + """Return free glibc arena pages to Linux without making it a dependency.""" + import sys + + if not sys.platform.startswith("linux"): + return False + try: + import ctypes + + malloc_trim = ctypes.CDLL(None).malloc_trim + malloc_trim.argtypes = [ctypes.c_size_t] + malloc_trim.restype = ctypes.c_int + return bool(malloc_trim(0)) + except (AttributeError, OSError): + return False + + +def reclaim_idle_browser_memory() -> bool: + """Collect closed Playwright cycles and trim idle allocator arenas once.""" + global _browser_reclaimed_generation + + if not _browser_memory_reclaim_lock.acquire(blocking=False): + return False + try: + with _browser_semaphore_lock: + target_generation = _browser_shutdown_generation + if ( + _browser_active != 0 + or _browser_waiters != 0 + or target_generation <= _browser_reclaimed_generation + ): + return False + with _retained_browser_resources_lock: + if _retained_browser_resources: + return False + with _tabs_lock: + if _tabs: + return False + + # Browser/context/page handles have left their owner-thread stacks by + # this boundary. CloakBrowser and Playwright create strong reference + # cycles, so refcounting alone cannot release their response buffers. + import gc + + collected = gc.collect() + trimmed = _trim_process_heap() + with _browser_semaphore_lock: + _browser_reclaimed_generation = max( + _browser_reclaimed_generation, + target_generation, + ) + logger.debug( + "browser_engine: reclaimed idle browser memory " + "(generation=%d, collected=%d, trimmed=%s)", + target_generation, + collected, + trimmed, + ) + return True + except Exception as exc: + logger.debug("browser_engine: idle memory reclamation failed: %s", exc) + return False + finally: + _browser_memory_reclaim_lock.release() def _build_browser_args(config: dict[str, Any] | None = None) -> list[str]: @@ -63,57 +569,159 @@ def _build_browser_args(config: dict[str, Any] | None = None) -> list[str]: return args +def _optional_boolean_state(obj: Any, name: str) -> bool | None: + try: + value = getattr(obj, name) + except (AttributeError, TypeError): + return None + except Exception: + return None + if callable(value): + try: + value = value() + except Exception: + return None + return value if isinstance(value, bool) else None + + +def _context_is_confirmed_closed(context: Any) -> bool: + """Return True only when the driver exposes an affirmative closed state.""" + for attribute in ("is_closed", "closed", "_closed"): + if _optional_boolean_state(context, attribute) is True: + return True + return False + + +def _browser_is_confirmed_closed(browser: Any) -> bool: + """Return True when the browser driver confirms its process is disconnected.""" + if _optional_boolean_state(browser, "is_connected") is False: + return True + return _context_is_confirmed_closed(browser) + + +def _close_resource_with_confirmation( + resource: Any, + *, + is_browser: bool = False, + attempts: int = 2, +) -> tuple[bool, list[Exception]]: + """Close a driver resource without treating an exception as successful cleanup.""" + errors: list[Exception] = [] + confirmed_closed = ( + _browser_is_confirmed_closed if is_browser else _context_is_confirmed_closed + ) + for attempt in range(max(1, attempts)): + try: + resource.close() + return True, errors + except Exception as exc: + errors.append(exc) + if confirmed_closed(resource): + return True, errors + if attempt + 1 < attempts: + time.sleep(0.05) + return False, errors + + +def _shared_browser_is_usable(browser: Any, context: Any) -> bool: + if browser is None or context is None: + return False + if _optional_boolean_state(browser, "is_connected") is False: + return False + for attribute in ("is_closed", "closed", "_closed"): + if _optional_boolean_state(context, attribute) is True: + return False + return True + + +def _thread_browser_resources_present() -> bool: + """Return whether this thread still owns shared-engine state to clean up.""" + return bool( + getattr(_tls, "browser", None) is not None + or getattr(_tls, "context", None) is not None + or getattr(_tls, "browser_lease", None) is not None + or getattr(_tls, "semaphore", None) is not None + or getattr(_tls, "tab_ids", None) + or _owner_retained_browser_resources() + ) + def _get_shared_browser(config: dict[str, Any] | None = None): """Get or create a browser for the current thread. Returns (browser, context).""" browser = getattr(_tls, "browser", None) context = getattr(_tls, "context", None) - if browser is not None: + if _shared_browser_is_usable(browser, context): return browser, context - - # Playwright's Sync API cannot run inside a running asyncio event loop. - # Detect that case and bail out early with a clear message instead of - # hitting a confusing "asyncio.run() cannot be called from a running - # event loop" deep inside CloakBrowser. + if ( + browser is not None + or context is not None + or getattr(_tls, "browser_lease", None) is not None + or getattr(_tls, "semaphore", None) is not None + or getattr(_tls, "tab_ids", None) + ): + logger.warning("browser_engine: replacing stale thread-local browser") + if not shutdown_shared_browser(): + raise RuntimeError( + "browser_engine: stale thread-local browser could not be closed" + ) + + # Playwright Sync API cannot run inside an asyncio event loop try: import asyncio asyncio.get_running_loop() raise RuntimeError( - "CloakBrowser (Playwright Sync API) cannot run inside an asyncio " - "event loop. Use the HTTP download sources instead, or call from " - "a synchronous context." + "CloakBrowser (Playwright Sync API) cannot run inside an asyncio event loop. " + "Use the HTTP download sources instead, or run outside of async context." ) except RuntimeError as e: if "cannot run inside" in str(e): raise - # No running loop — safe to proceed. + # No running loop — OK to proceed + pass if not _check_cloakbrowser(): raise RuntimeError("cloakbrowser not installed. Run: pip install cloakbrowser") - # Platform compat shim - try: - from ..institutional.cloakbrowser_compat import ensure_cloakbrowser_platform_compatible - ensure_cloakbrowser_platform_compatible() - except Exception: - try: - from .institutional.cloakbrowser_compat import ensure_cloakbrowser_platform_compatible - ensure_cloakbrowser_platform_compatible() - except Exception: - pass - from cloakbrowser import launch + from .cloakbrowser_compat import launch_with_driver_cleanup - headless = False + # Auto-detect headless: Docker/CI environments have no DISPLAY + import os + _has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")) + headless = not _has_display # default: headless when no display humanize = True if config: - headless = config.get("browser_headless", False) + if _has_display: + # Only respect config override when a display is available + headless = config.get("browser_headless", False) humanize = config.get("browser_humanize", True) args = _build_browser_args(config) - browser = launch(headless=headless, humanize=humanize, args=args) - context = browser.new_context() + lease = _retain_browser_slot(config) + browser = None + context = None + try: + browser = launch_with_driver_cleanup( + launch, + headless=headless, + humanize=humanize, + args=args, + ) + context = browser.new_context() + _seed_context_cookies(context, config) + except Exception: + if browser is None: + lease.close() + else: + _tls.browser = browser + _tls.context = context + _tls.browser_lease = lease + _tls.semaphore = lease.semaphore + shutdown_shared_browser() + raise _tls.browser = browser _tls.context = context + _tls.browser_lease = lease + _tls.semaphore = lease.semaphore logger.info(f"browser_engine: browser ready for thread {_threading.current_thread().name}") return browser, context @@ -178,30 +786,150 @@ def get_browser_page(config: dict[str, Any] | None = None): return None -def shutdown_shared_browser(): +def get_persistent_context( + profile_dir: str | Path, + config: dict[str, Any] | None = None, +): + """Get or create a persistent browser context for fingerprint consistency. + + Unlike launch() + cookie restore, persistent context preserves: + - Browser fingerprint (canvas, WebGL, audio, fonts) + - Cookies and localStorage across restarts + - Login sessions without re-authentication + + This is the recommended approach for publisher sessions that need + stable identity across multiple download runs. + """ + if not _check_cloakbrowser(): + raise RuntimeError("cloakbrowser not installed. Run: pip install cloakbrowser") + + from cloakbrowser import launch_persistent_context + from .cloakbrowser_compat import launch_with_driver_cleanup + + import os + _has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")) + headless = not _has_display + humanize = True + if config: + if _has_display: + headless = config.get("browser_headless", False) + humanize = config.get("browser_humanize", True) + + args = _build_browser_args(config) + profile_path = Path(profile_dir) + profile_path.mkdir(parents=True, exist_ok=True) + + lease = _retain_browser_slot(config) + ctx = None + try: + ctx = launch_with_driver_cleanup( + launch_persistent_context, + str(profile_path), + headless=headless, + humanize=humanize, + args=args, + ) + _seed_context_cookies(ctx, config) + except Exception: + context_closed = ctx is None + if ctx is not None: + context_closed, close_errors = _close_resource_with_confirmation(ctx) + if not context_closed: + retained_context = _LeasedPersistentContext(ctx, lease) + _register_retained_browser_resource(retained_context) + logger.warning( + "browser_engine: failed to close persistent context after " + "initialization error; retaining its owner-thread handle " + "and browser permit: " + + "; ".join(str(error) for error in close_errors) + ) + if context_closed: + lease.close() + raise + logger.info(f"browser_engine: persistent context ready at {profile_path}") + return _LeasedPersistentContext(ctx, lease) + + +def shutdown_shared_browser() -> bool: """Shut down the current thread's browser. Call on thread exit or process exit.""" + context = getattr(_tls, "context", None) browser = getattr(_tls, "browser", None) + lease = getattr(_tls, "browser_lease", None) + semaphore = getattr(_tls, "semaphore", None) + if ( + context is None + and browser is None + and lease is None + and semaphore is None + and not getattr(_tls, "tab_ids", None) + and not _owner_retained_browser_resources() + ): + return True + + tab_close_errors = _close_thread_tabs() + context_closed = context is None + browser_closed = browser is None + context_errors: list[Exception] = [] + browser_errors: list[Exception] = [] + + if context is not None: + context_closed, context_errors = _close_resource_with_confirmation(context) if browser is not None: - try: - browser.close() - except Exception: - pass - _tls.browser = None - _tls.context = None + browser_closed, browser_errors = _close_resource_with_confirmation( + browser, + is_browser=True, + ) + if browser_closed: + # A normally closed/disconnected browser owns and terminates all of + # its contexts, even if an earlier context.close() call raised. + context_closed = True + + _tls.context = None if context_closed else context + _tls.browser = None if browser_closed else browser + shared_shutdown_complete = context_closed and browser_closed + if shared_shutdown_complete: + _tls.browser_lease = None + _tls.semaphore = None + if lease is not None: + lease.close() + elif semaphore is not None: + _release_browser_slot(semaphore) + else: + # Keep both the permit and every still-live driver handle. A later + # shutdown call can retry without allowing a replacement Chromium to + # exceed the configured browser limit. + _tls.browser_lease = lease + _tls.semaphore = semaphore + + retained_closed, retained_errors = _retry_retained_browser_resources() + shutdown_complete = shared_shutdown_complete and retained_closed + + close_errors = list(tab_close_errors) + if not context_closed: + close_errors.append( + "context: " + "; retry: ".join(str(error) for error in context_errors) + ) + if not browser_closed: + close_errors.append( + "browser: " + "; retry: ".join(str(error) for error in browser_errors) + ) + if not retained_closed: + close_errors.append( + "retained resource: " + + "; retry: ".join(str(error) for error in retained_errors) + ) + if close_errors: + logger.warning( + "browser_engine: browser shutdown incomplete: " + "; ".join(close_errors) + ) + else: logger.info("browser_engine: browser shut down") + return shutdown_complete -def _ensure_compat(): - """Ensure CloakBrowser platform compatibility.""" - try: - from ..institutional.cloakbrowser_compat import ensure_cloakbrowser_platform_compatible - ensure_cloakbrowser_platform_compatible() - except Exception: - try: - from .institutional.cloakbrowser_compat import ensure_cloakbrowser_platform_compatible - ensure_cloakbrowser_platform_compatible() - except Exception: - pass +def _ensure_compat() -> bool: + """Backward-compatible alias for the real, package-local runtime setup.""" + return _prepare_cloakbrowser_runtime() # --------------------------------------------------------------------------- @@ -219,31 +947,170 @@ def is_available(config: dict[str, Any] | None = None) -> bool: # tab-based workflows within a single operation (thread-safe: one thread per tab). # --------------------------------------------------------------------------- -_tabs: dict[str, Any] = {} # tab_id → page -_captured: dict[str, list] = {} # tab_id → captured PDF responses +_tabs: dict[str, Any] = {} # tab_id -> page +_tab_owners: dict[str, int] = {} +_captured: dict[str, list] = {} # tab_id -> captured PDF responses +_tabs_lock = _threading.Lock() + +_CookieKey = tuple[str, str, str] +_imported_cookies: dict[Path, dict[_CookieKey, dict[str, Any]]] = {} +_imported_cookies_lock = _threading.Lock() +_imported_cookie_path_locks: dict[Path, Any] = {} +_imported_cookie_path_locks_lock = _threading.Lock() + + +def _is_valid_imported_cookie(cookie: dict[str, Any]) -> bool: + return bool( + cookie.get("name") + and (cookie.get("domain") or cookie.get("url")) + and cookie.get("path", "/").startswith("/") + ) + + +def _persistent_cookie_path(config: dict[str, Any] | None) -> Path: + from .config import DATA_DIR + + cache_dir = Path( + (config or {}).get("cache_dir") or str(DATA_DIR / "cache") + ).expanduser() + return (cache_dir / "imported_browser_cookies.json").resolve() + + +def _imported_cookie_path_lock(path: Path): + with _imported_cookie_path_locks_lock: + lock = _imported_cookie_path_locks.get(path) + if lock is None: + lock = _threading.RLock() + _imported_cookie_path_locks[path] = lock + return lock + + +def _load_persisted_imported_cookies( + config: dict[str, Any] | None, +) -> list[dict[str, Any]]: + path = _persistent_cookie_path(config) + with _imported_cookie_path_lock(path): + if not path.exists(): + return [] + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except Exception as exc: + logger.warning(f"browser_engine: failed to read imported cookies: {exc}") + return [] + if not isinstance(payload, list): + return [] + return [ + dict(cookie) + for cookie in payload + if isinstance(cookie, dict) and _is_valid_imported_cookie(cookie) + ] + + +def _cookie_key(cookie: dict[str, Any]) -> _CookieKey: + return ( + str(cookie.get("domain", "")), + str(cookie.get("path", "/")), + str(cookie.get("name", "")), + ) + + +def _remember_imported_cookies( + config: dict[str, Any] | None, + cookies: list[dict[str, Any]], +) -> None: + path = _persistent_cookie_path(config) + with _imported_cookie_path_lock(path): + with _imported_cookies_lock: + cookie_store = _imported_cookies.setdefault(path, {}) + for cookie in cookies: + if not _is_valid_imported_cookie(cookie): + continue + cookie_store[_cookie_key(cookie)] = dict(cookie) + + +def _persist_imported_cookies(config: dict[str, Any] | None) -> bool: + path = _persistent_cookie_path(config) + temp_path = path.with_name(f"{path.name}.{uuid.uuid4().hex}.tmp") + with _imported_cookie_path_lock(path): + with _imported_cookies_lock: + cookies = [ + dict(cookie) + for cookie in _imported_cookies.get(path, {}).values() + ] + try: + path.parent.mkdir(parents=True, exist_ok=True) + temp_path.write_text( + json.dumps(cookies, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + temp_path.replace(path) + return True + except Exception as exc: + logger.warning(f"browser_engine: failed to persist imported cookies: {exc}") + return False + finally: + try: + temp_path.unlink(missing_ok=True) + except OSError: + pass + + +def _seed_context_cookies(context: Any, config: dict[str, Any] | None) -> None: + path = _persistent_cookie_path(config) + with _imported_cookie_path_lock(path): + _remember_imported_cookies( + config, + _load_persisted_imported_cookies(config), + ) + with _imported_cookies_lock: + cookie_items = [ + (key, dict(cookie)) + for key, cookie in _imported_cookies.get(path, {}).items() + ] + rejected = [] + for key, cookie in cookie_items: + try: + context.add_cookies([cookie]) + except Exception as exc: + rejected.append((key, cookie)) + logger.warning( + "browser_engine: rejected imported cookie " + f"{cookie.get('name', '')}: {exc}" + ) + if rejected: + with _imported_cookies_lock: + cookie_store = _imported_cookies.get(path, {}) + for key, rejected_cookie in rejected: + if cookie_store.get(key) == rejected_cookie: + cookie_store.pop(key, None) + _persist_imported_cookies(config) def _register_tab(browser, context, page) -> str: tab_id = uuid.uuid4().hex[:12] - _tabs[tab_id] = page - _captured[tab_id] = [] + with _tabs_lock: + _tabs[tab_id] = page + _tab_owners[tab_id] = _threading.get_ident() + _captured[tab_id] = [] + tab_ids = getattr(_tls, "tab_ids", None) + if tab_ids is None: + tab_ids = set() + _tls.tab_ids = tab_ids + tab_ids.add(tab_id) - # Listen for PDF responses + # Record PDF response metadata without retaining response bodies in memory. def _on_response(response): try: ct = response.headers.get("content-type", "") if "pdf" in ct or "octet-stream" in ct: - try: - body = response.body() - if body[:5] == b"%PDF-": - _captured[tab_id].append({ + with _tabs_lock: + captured = _captured.get(tab_id) + if captured is not None: + captured.append({ "url": response.url, "status": response.status, "contentType": ct, - "dataBase64": base64.b64encode(body).decode(), }) - except Exception: - pass except Exception: pass @@ -257,10 +1124,33 @@ def _on_response(response): def _resolve_tab(tab_id: str): """Look up page for a tab_id. Returns page or None.""" - page = _tabs.get(tab_id) + with _tabs_lock: + owner = _tab_owners.get(tab_id) + page = _tabs.get(tab_id) + if owner is not None and owner != _threading.get_ident(): + logger.warning(f"browser_engine: tab {tab_id} belongs to another thread") + return None return page +def _close_thread_tabs() -> list[str]: + """Close and forget pages owned by the current Playwright thread.""" + tab_ids = tuple(getattr(_tls, "tab_ids", ())) + _tls.tab_ids = set() + errors: list[str] = [] + for tab_id in tab_ids: + with _tabs_lock: + page = _tabs.pop(tab_id, None) + _tab_owners.pop(tab_id, None) + _captured.pop(tab_id, None) + if page is not None: + try: + page.close() + except Exception as exc: + errors.append(f"tab {tab_id}: {exc}") + return errors + + def solve_url( url: str, config: dict[str, Any], @@ -336,7 +1226,7 @@ def get_html( def import_cookies(cookie_file: str | Path, config: dict[str, Any], *, domain_suffix: str | None = None) -> int: - """Import Netscape-format cookies into the shared browser context. Returns count imported.""" + """Remember cookies and add them to the current thread's context, if any.""" try: text = Path(cookie_file).read_text(encoding="utf-8") except Exception as e: @@ -347,14 +1237,34 @@ def import_cookies(cookie_file: str | Path, config: dict[str, Any], *, domain_su return 0 if domain_suffix: cookies = [c for c in cookies if domain_suffix in c.get("domain", "")] - try: - _, ctx = _get_shared_browser(config) - ctx.add_cookies(cookies) - logger.info(f"browser_engine: imported {len(cookies)} cookies") - return len(cookies) - except Exception as e: - logger.info(f"browser_engine: import_cookies error: {e}") + cookies = [cookie for cookie in cookies if _is_valid_imported_cookie(cookie)] + if not cookies: + return 0 + + path = _persistent_cookie_path(config) + with _imported_cookie_path_lock(path): + _remember_imported_cookies( + config, + _load_persisted_imported_cookies(config), + ) + _remember_imported_cookies(config, cookies) + persisted = _persist_imported_cookies(config) + + ctx = getattr(_tls, "context", None) + if ctx is not None: + try: + ctx.add_cookies(cookies) + except Exception as e: + logger.info( + "browser_engine: import_cookies failed on the current context; " + f"retiring the thread-local browser: {e}" + ) + shutdown_shared_browser() + return len(cookies) if persisted else 0 + if ctx is None and not persisted: return 0 + logger.info(f"browser_engine: imported {len(cookies)} cookies") + return len(cookies) def evaluate_js( @@ -378,6 +1288,7 @@ def evaluate_js( def create_tab(url: str, config: dict[str, Any], *, timeout: float = 30.0) -> str | None: """Create a new tab (page) in the shared browser and navigate to URL. Returns tab_id or None.""" + page = None try: browser, context = _get_shared_browser(config) page = context.new_page() @@ -386,19 +1297,51 @@ def create_tab(url: str, config: dict[str, Any], *, timeout: float = 30.0) -> st return tab_id except Exception as e: logger.info(f"browser_engine: create_tab failed - {e}") + if page is not None: + try: + page.close() + except Exception: + pass + shutdown_shared_browser() return None def close_tab(tab_id: str, config: dict[str, Any]) -> None: """Close a browser tab (page only, not the shared browser).""" - page = _resolve_tab(tab_id) - if page: + with _tabs_lock: + owner = _tab_owners.get(tab_id) + if owner is not None and owner != _threading.get_ident(): + logger.warning(f"browser_engine: refusing cross-thread close for tab {tab_id}") + return + page = _tabs.get(tab_id) + if page is None: + return + + close_error = None + for _attempt in range(2): try: page.close() - except Exception: - pass - _tabs.pop(tab_id, None) - _captured.pop(tab_id, None) + close_error = None + break + except Exception as exc: + close_error = exc + + if close_error is not None: + logger.warning( + f"browser_engine: failed to close tab {tab_id} after retry; " + f"retiring the thread-local browser: {close_error}" + ) + shutdown_shared_browser() + return + + with _tabs_lock: + if _tabs.get(tab_id) is page: + _tabs.pop(tab_id, None) + _tab_owners.pop(tab_id, None) + _captured.pop(tab_id, None) + tab_ids = getattr(_tls, "tab_ids", None) + if tab_ids is not None: + tab_ids.discard(tab_id) def navigate_tab(tab_id: str, url: str, config: dict[str, Any], *, timeout: float = 30.0) -> bool: @@ -434,6 +1377,7 @@ def download_pdf_via_browser( config: dict[str, Any], *, timeout: float = 60.0, + cancel_event: _threading.Event | None = None, ) -> bool: """Download a PDF URL via CloakBrowser shared browser. Returns True on success. @@ -443,37 +1387,112 @@ def download_pdf_via_browser( 3. PDF link discovery in page DOM 4. Download button click """ + context = None page = None + capture_lock = _threading.Lock() + captured_response: dict[str, Any] = { + "body": None, + "in_progress": False, + "accepting": True, + } + cancel_event = _effective_cancel_event(cancel_event) try: + if _cancelled(cancel_event): + return False _, context = _get_shared_browser(config) + if _cancelled(cancel_event): + return False page = context.new_page() - # Set up response listener for PDF captures - captured_responses: list[dict] = [] - def _on_response(response): + with capture_lock: + if ( + not captured_response["accepting"] + or _cancelled(cancel_event) + or captured_response["body"] is not None + or captured_response["in_progress"] + ): + return + captured_response["in_progress"] = True try: ct = response.headers.get("content-type", "") if "pdf" in ct or "octet-stream" in ct: - try: - body = response.body() - if body[:5] == b"%PDF-": - captured_responses.append({ - "url": response.url, - "dataBase64": base64.b64encode(body).decode(), - }) - except Exception: - pass + body = response.body() + if body[:5] == b"%PDF-" and not _cancelled(cancel_event): + with capture_lock: + if ( + captured_response["accepting"] + and captured_response["body"] is None + ): + captured_response["body"] = body except Exception: pass + finally: + with capture_lock: + captured_response["in_progress"] = False try: page.on("response", _on_response) except Exception: pass - page.goto(pdf_url, wait_until="domcontentloaded", timeout=int(timeout * 1000)) - time.sleep(3) + page_loaded = False + if _is_pdf_url(pdf_url): + from .pdf_utils import is_pdf_file, publish_pdf_file_atomic + + download_path = output_path.with_name( + f".{output_path.name}.{uuid.uuid4().hex}.download" + ) + download_observed = False + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + with page.expect_download( + timeout=int(min(timeout, 30.0) * 1000) + ) as download_info: + try: + page.goto( + pdf_url, + wait_until="commit", + timeout=int(timeout * 1000), + ) + page_loaded = True + except Exception as exc: + if "Download is starting" not in str(exc): + raise + download = download_info.value + download_observed = True + if _cancelled(cancel_event): + with contextlib.suppress(Exception): + download.cancel() + return False + download.save_as(str(download_path)) + if ( + is_pdf_file(download_path) + and publish_pdf_file_atomic( + download_path, + output_path, + cancel_event, + ) + ): + logger.info( + "browser_engine: downloaded PDF via browser download event" + ) + return True + except Exception as exc: + logger.info( + "browser_engine: direct download event unavailable: %s", + exc, + ) + finally: + with contextlib.suppress(OSError): + download_path.unlink(missing_ok=True) + if download_observed: + return False + + if not page_loaded: + page.goto(pdf_url, wait_until="domcontentloaded", timeout=int(timeout * 1000)) + if _wait_or_cancel(cancel_event, 3): + return False # Check for anti-bot challenges html = "" @@ -491,20 +1510,22 @@ def _on_response(response): "altcha", "你是机器人吗", "not a robot", "nope", ]): logger.info("browser_engine: anti-bot challenge detected, waiting...") - time.sleep(10) + if _wait_or_cancel(cancel_event, 10): + return False current_url = page.url + if _cancelled(cancel_event): + return False # Strategy 0: Network response capture - for resp in captured_responses: - data = resp.get("dataBase64", "") - if data: - pdf_bytes = base64.b64decode(data) - if pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - logger.info(f"browser_engine: downloaded {len(pdf_bytes)} bytes via network capture") - return True + with capture_lock: + pdf_bytes = captured_response["body"] + captured_response["body"] = None + if pdf_bytes is not None and len(pdf_bytes) > 5000: + if _write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + logger.info(f"browser_engine: downloaded {len(pdf_bytes)} bytes via network capture") + return True + return False # Build candidate fetch paths fetch_paths: list[str] = [] @@ -531,6 +1552,8 @@ def _on_response(response): # Strategy 1: In-browser fetch API for fetch_path in fetch_paths: + if _cancelled(cancel_event): + return False logger.info(f"browser_engine: trying in-browser fetch {origin}{fetch_path[:60]}") try: pdf_b64 = page.evaluate(f""" @@ -554,15 +1577,17 @@ def _on_response(response): }} }})() """) + if _cancelled(cancel_event): + return False if isinstance(pdf_b64, str) and pdf_b64.startswith("data:"): header, data = pdf_b64.split(",", 1) pdf_bytes = base64.b64decode(data) if pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - logger.info(f"browser_engine: downloaded {len(pdf_bytes)} bytes via in-browser fetch") - return True + if _write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + logger.info(f"browser_engine: downloaded {len(pdf_bytes)} bytes via in-browser fetch") + return True + return False else: logger.info(f"browser_engine: fetch returned non-PDF ({len(pdf_bytes)} bytes)") else: @@ -571,6 +1596,8 @@ def _on_response(response): logger.info(f"browser_engine: fetch error: {e}") # Strategy 2: PDF link discovery in DOM + if _cancelled(cancel_event): + return False try: pdf_link = page.evaluate(""" (() => { @@ -594,6 +1621,8 @@ def _on_response(response): return null; })() """) + if _cancelled(cancel_event): + return False if pdf_link and isinstance(pdf_link, str) and pdf_link.startswith("http"): logger.info(f"browser_engine: found PDF link: {pdf_link[:80]}") @@ -622,21 +1651,25 @@ def _on_response(response): }} }})() """) + if _cancelled(cancel_event): + return False if isinstance(pdf_b64, str) and pdf_b64.startswith("data:"): header, data = pdf_b64.split(",", 1) pdf_bytes = base64.b64decode(data) if pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - logger.info(f"browser_engine: downloaded {len(pdf_bytes)} bytes via PDF link fetch") - return True + if _write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + logger.info(f"browser_engine: downloaded {len(pdf_bytes)} bytes via PDF link fetch") + return True + return False except Exception as e: logger.info(f"browser_engine: PDF link fetch error: {e}") except Exception as e: logger.info(f"browser_engine: DOM scan error: {e}") # Strategy 3: Click download button + if _cancelled(cancel_event): + return False try: clicked = page.evaluate(""" (() => { @@ -648,22 +1681,44 @@ def _on_response(response): return false; })() """) + if _cancelled(cancel_event): + return False if clicked: logger.info("browser_engine: clicked download button, waiting...") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return False except Exception as e: logger.info(f"browser_engine: click error: {e}") return False except Exception as e: logger.info(f"browser_engine: download_pdf_via_browser error: {e}") + if context is not None and page is None: + shutdown_shared_browser() return False finally: - if page: + with capture_lock: + captured_response["accepting"] = False + captured_response["body"] = None + if page is not None: try: - page.close() + page.remove_listener("response", _on_response) except Exception: pass + close_error = None + for _attempt in range(2): + try: + page.close() + close_error = None + break + except Exception as exc: + close_error = exc + if close_error is not None: + logger.warning( + "browser_engine: failed to close PDF download page after retry; " + f"retiring the thread-local browser: {close_error}" + ) + shutdown_shared_browser() # Backward-compat alias @@ -675,6 +1730,7 @@ def _is_pdf_url(url: str) -> bool: lower = url.lower() return ( lower.endswith(".pdf") + or urlparse(lower).path.endswith("/pdf") or "/pdf/" in lower or "content/pdf" in lower or "pdfdirect" in lower @@ -696,28 +1752,19 @@ def fetch_url( logger.info(f"browser_engine: fetch_url - tab {tab_id} not found") return None + response = None try: - page.goto(url, wait_until="domcontentloaded", timeout=int(timeout * 1000)) + response = page.goto(url, wait_until="domcontentloaded", timeout=int(timeout * 1000)) except Exception as e: logger.info(f"browser_engine: fetch_url navigate error: {e}") - # Check captured responses - for resp in _captured.get(tab_id, []): - data = resp.get("dataBase64", "") - if data: - pdf_bytes = base64.b64decode(data) - if pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: - return {"status": "ok", "bytes": len(pdf_bytes), "data": pdf_bytes} - - # Also check direct response body - try: - response = page.goto(url, wait_until="commit", timeout=int(timeout * 1000)) - if response is not None: + if response is not None: + try: body = response.body() if body[:5] == b"%PDF-" and len(body) > 5000: return {"status": "ok", "bytes": len(body), "data": body} - except Exception: - pass + except Exception: + pass return None @@ -729,16 +1776,24 @@ def get_captured_responses( consume: bool = True, ) -> list[dict[str, Any]]: """Get captured PDF responses for a tab. Optionally consume (clear) them.""" - captured = _captured.get(tab_id, []) - result = list(captured) - if consume: - _captured[tab_id] = [] + with _tabs_lock: + captured = _captured.get(tab_id, []) + result = list(captured) + if consume and tab_id in _captured: + _captured[tab_id] = [] return result def close_all_tabs(config: dict[str, Any]) -> None: - """Close all tracked browser tabs (pages only, not the shared browser).""" - for tab_id in list(_tabs.keys()): + """Close tabs owned by the current Playwright thread.""" + owner = _threading.get_ident() + with _tabs_lock: + tab_ids = [ + tab_id + for tab_id, tab_owner in _tab_owners.items() + if tab_owner == owner + ] + for tab_id in tab_ids: close_tab(tab_id, config) diff --git a/src/scansci_pdf/browser_login.py b/src/scansci_pdf/browser_login.py index fedc828..35ed262 100644 --- a/src/scansci_pdf/browser_login.py +++ b/src/scansci_pdf/browser_login.py @@ -5,13 +5,17 @@ import json import time import atexit +import threading from pathlib import Path from typing import Any +from .cloakbrowser_compat import launch_with_driver_cleanup, prepare_cloakbrowser_runtime + try: + prepare_cloakbrowser_runtime() from cloakbrowser import launch _HAS_CLOAKBROWSER = True -except ImportError: +except Exception: launch = None # type: ignore[assignment] _HAS_CLOAKBROWSER = False from .log import get_logger @@ -31,42 +35,114 @@ def __init__(self): self._context = None self._page = None self._cookies_saved = False - - @property - def is_alive(self) -> bool: + self._owner_thread: threading.Thread | None = None + self._owner_thread_id: int | None = None + self._browser_lease = None + self._lifecycle_lock = threading.RLock() + + def _assert_owner_locked(self, operation: str) -> None: + owner = self._owner_thread + current = threading.get_ident() + if owner is not None and owner is not threading.current_thread(): + raise RuntimeError( + "PersistentBrowser is owned by another thread; " + f"cannot {operation} Playwright resources from thread {current} " + f"(owner={self._owner_thread_id})" + ) + + def _is_alive_locked(self) -> bool: if self._browser is None: return False try: self._page.url # noqa: B018 return True except Exception: - self._cleanup() + if not self._cleanup_locked(): + raise RuntimeError( + "PersistentBrowser resources became stale and could not be closed" + ) return False + @property + def is_alive(self) -> bool: + with self._lifecycle_lock: + self._assert_owner_locked("inspect") + return self._is_alive_locked() + def get_page(self, config: dict[str, Any] | None = None): """Get or create the browser page. Returns (context, page).""" - if self.is_alive: - return self._context, self._page - return self._start(config) + with self._lifecycle_lock: + self._assert_owner_locked("reuse") + if self._is_alive_locked(): + return self._context, self._page + return self._start_locked(config) def _start(self, config: dict[str, Any] | None = None): """Start a new browser instance. Restores saved state if available.""" + with self._lifecycle_lock: + self._assert_owner_locked("start") + if self._is_alive_locked(): + return self._context, self._page + return self._start_locked(config) + + def _start_locked(self, config: dict[str, Any] | None = None): + """Start resources while the lifecycle lock is held by their owner.""" if not _HAS_CLOAKBROWSER: raise RuntimeError("cloakbrowser not installed. Run: pip install cloakbrowser") log.info(" [browser] Starting persistent browser...") - self._browser = launch( - headless=False, humanize=True, - args=["--disable-features=CrossOriginOpenerPolicy"], - ) - self._context = self._browser.new_context() - self._page = self._context.new_page() - - if config: - self._restore_state(config) - - return self._context, self._page - - def _restore_state(self, config: dict[str, Any]): + browser = None + context = None + page = None + slot_lease = None + try: + from . import browser_engine + slot_lease = browser_engine._retain_browser_slot(config) + browser = launch_with_driver_cleanup( + launch, + headless=False, humanize=True, + args=["--disable-features=CrossOriginOpenerPolicy"], + ) + context = browser.new_context() + page = context.new_page() + if config: + self._restore_state(config, context=context, page=page) + except Exception: + resources_closed = self._close_resources(page, context, browser) + if resources_closed: + if slot_lease is not None: + slot_lease.close() + self._browser = None + self._context = None + self._page = None + self._browser_lease = None + self._owner_thread = None + self._owner_thread_id = None + else: + # Preserve the exact handles and permit for an owner-thread + # retry. Launching a replacement here could exceed the cap. + self._browser = browser + self._context = context + self._page = page + self._browser_lease = slot_lease + self._owner_thread = threading.current_thread() + self._owner_thread_id = threading.get_ident() + raise + + self._browser = browser + self._context = context + self._page = page + self._browser_lease = slot_lease + self._owner_thread = threading.current_thread() + self._owner_thread_id = threading.get_ident() + return context, page + + def _restore_state( + self, + config: dict[str, Any], + *, + context=None, + page=None, + ): """Restore saved cookies and localStorage into the browser.""" from .config import DATA_DIR cache_dir = Path(config.get("cache_dir", str(DATA_DIR / "cache"))) @@ -78,21 +154,31 @@ def _restore_state(self, config: dict[str, Any]): except Exception: log.info(" [browser] browser_state.json corrupted, starting fresh") return + if not isinstance(state, dict): + log.info(" [browser] browser_state.json has invalid structure, starting fresh") + return + + if context is None: + context = self._context + if page is None: + page = self._page cookies = state.get("cookies", []) - if cookies: + if cookies and context is not None: try: - self._context.add_cookies(cookies) + context.add_cookies(cookies) log.info(f" [browser] Restored {len(cookies)} cookies") except Exception as e: log.info(f" [browser] Cookie restore warning: {e}") storage = state.get("localStorage", {}) + if not isinstance(storage, dict) or page is None: + storage = {} for origin, items in storage.items(): try: - self._page.goto(origin, wait_until="commit", timeout=10000) + page.goto(origin, wait_until="commit", timeout=10000) for key, value in items.items(): - self._page.evaluate(f"localStorage.setItem({json.dumps(key)}, {json.dumps(value)})") + page.evaluate(f"localStorage.setItem({json.dumps(key)}, {json.dumps(value)})") except Exception as e: log.info(f" [browser] localStorage restore failed for {origin}: {e}") @@ -100,77 +186,137 @@ def _restore_state(self, config: dict[str, Any]): def save_cookies(self, config: dict[str, Any]): """Save current browser state (cookies + localStorage) to disk.""" - if not self._context: - return - try: - from .config import DATA_DIR - cache_dir = Path(config.get("cache_dir", str(DATA_DIR / "cache"))) - cache_dir.mkdir(parents=True, exist_ok=True) - - cookies = self._context.cookies() - - localStorage = {} - for page in self._context.pages: - try: - url = page.url - if url.startswith("http"): - from urllib.parse import urlparse - origin = f"{urlparse(url).scheme}://{urlparse(url).hostname}" - items = page.evaluate(""" - (() => { - const items = {}; - for (let i = 0; i < localStorage.length; i++) { - const key = localStorage.key(i); - items[key] = localStorage.getItem(key); - } - return items; - })() - """) - if items: - localStorage[origin] = items - except Exception: - pass - - state = {"cookies": cookies, "localStorage": localStorage} - state_file = cache_dir / "browser_state.json" - state_file.write_text(json.dumps(state, indent=2, ensure_ascii=False), encoding="utf-8") - - cookie_file = cache_dir / "instsci-cookies.json" - cookie_data = [ - {"name": c["name"], "value": c["value"], "domain": c.get("domain", ""), "path": c.get("path", "/")} - for c in cookies - ] - cookie_file.write_text(json.dumps(cookie_data, indent=2, ensure_ascii=False), encoding="utf-8") - - netscape_file = cache_dir / "instsci-cookies.txt" - from .browser_cookies import cookies_to_netscape - netscape_file.write_text(cookies_to_netscape(cookies), encoding="utf-8") - - self._cookies_saved = True - log.info(f" [browser] Saved {len(cookies)} cookies + {len(localStorage)} localStorage origins") - except Exception as e: - log.info(f" [browser] Failed to save state: {e}") + with self._lifecycle_lock: + self._assert_owner_locked("save") + if not self._context: + return + try: + from .config import DATA_DIR + cache_dir = Path(config.get("cache_dir", str(DATA_DIR / "cache"))) + cache_dir.mkdir(parents=True, exist_ok=True) - def _cleanup(self): - """Close browser gracefully.""" - try: - if self._browser: - self._browser.close() - except Exception: - pass + cookies = self._context.cookies() + + localStorage = {} + for page in self._context.pages: + try: + url = page.url + if url.startswith("http"): + from urllib.parse import urlparse + origin = f"{urlparse(url).scheme}://{urlparse(url).hostname}" + items = page.evaluate(""" + (() => { + const items = {}; + for (let i = 0; i < localStorage.length; i++) { + const key = localStorage.key(i); + items[key] = localStorage.getItem(key); + } + return items; + })() + """) + if items: + localStorage[origin] = items + except Exception: + pass + + state = {"cookies": cookies, "localStorage": localStorage} + state_file = cache_dir / "browser_state.json" + state_file.write_text(json.dumps(state, indent=2, ensure_ascii=False), encoding="utf-8") + + cookie_file = cache_dir / "instsci-cookies.json" + cookie_data = [ + {"name": c["name"], "value": c["value"], "domain": c.get("domain", ""), "path": c.get("path", "/")} + for c in cookies + ] + cookie_file.write_text(json.dumps(cookie_data, indent=2, ensure_ascii=False), encoding="utf-8") + + netscape_file = cache_dir / "instsci-cookies.txt" + from .browser_cookies import cookies_to_netscape + netscape_file.write_text(cookies_to_netscape(cookies), encoding="utf-8") + + self._cookies_saved = True + log.info(f" [browser] Saved {len(cookies)} cookies + {len(localStorage)} localStorage origins") + except Exception as e: + log.info(f" [browser] Failed to save state: {e}") + + def _close_resources(self, page, context, browser) -> bool: + from . import browser_engine + + closed: dict[str, bool] = { + "page": page is None, + "context": context is None, + "browser": browser is None, + } + for label, resource in ( + ("page", page), + ("context", context), + ("browser", browser), + ): + if resource is None: + continue + resource_closed, errors = browser_engine._close_resource_with_confirmation( + resource, + is_browser=label == "browser", + ) + closed[label] = resource_closed + if not resource_closed: + log.info( + f" [browser] Failed to close {label} after retry: " + + "; ".join(str(error) for error in errors) + ) + + if browser is not None and closed["browser"]: + return True + if context is not None and closed["context"]: + return browser is None + return browser is None and context is None and closed["page"] + + def _cleanup_locked(self): + """Close resources while holding the lifecycle lock on the owner thread.""" + page = self._page + context = self._context + browser = self._browser + if not self._close_resources(page, context, browser): + return False + lease = self._browser_lease self._browser = None self._context = None self._page = None + self._browser_lease = None + self._owner_thread = None + self._owner_thread_id = None + if lease is not None: + lease.close() + return True + + def _cleanup(self): + """Close browser resources and detach ownership before calling drivers.""" + with self._lifecycle_lock: + self._assert_owner_locked("close") + if not self._cleanup_locked(): + raise RuntimeError( + "PersistentBrowser close could not be confirmed; " + "resources and global browser slot were retained" + ) def close(self): """Explicitly close the browser.""" self._cleanup() log.info(" [browser] Persistent browser closed") + def _close_at_exit(self) -> None: + """Best-effort process-exit cleanup without crossing Playwright threads.""" + try: + self._cleanup() + except RuntimeError: + # Logging handlers (including pytest's capture stream) may already + # be closed when atexit callbacks run. Cleanup must stay silent. + pass + # Module-level singleton _browser = PersistentBrowser() -atexit.register(_browser.close) +atexit.register(_browser._close_at_exit) def get_browser(config: dict[str, Any] | None = None): @@ -243,31 +389,41 @@ def open_login_browser( detect_login: Optional callable(browser_context, page) -> bool for custom login detection. max_wait: Max seconds to wait for login. auto_import: Whether to auto-import cookies into CloakBrowser. - keep_alive: If True, return (True, context, page) without closing browser. + keep_alive: If True, return (True, browser, context, page) without closing browser. publisher: Publisher name for remote assist display. Returns: - True if login succeeded, or (True, context, page) if keep_alive. + True if login succeeded, or (True, browser, context, page) if keep_alive. """ log.info(f" [browser] Opening stealth browser: {url}") print(f"\n 请在浏览器中登录 ({url})") print(" 程序会自动检测登录完成...\n") - # Start remote assist if port is configured - remote = None - if int(config.get("remote_assist_port", 0)) > 0: - from .remote_assist import RemoteAssist - remote = RemoteAssist(config, publisher=publisher) - remote.start() - remote.update_url(url) - if not _HAS_CLOAKBROWSER: log.info(" [browser] cloakbrowser not installed") return (False, None, None, None) if keep_alive else False + remote = None + browser = None + slot_lease = None + ownership_transferred = False try: - browser = launch(headless=False, humanize=True, - args=["--disable-features=CrossOriginOpenerPolicy"]) + if int(config.get("remote_assist_port", 0)) > 0: + from .remote_assist import RemoteAssist + remote = RemoteAssist(config, publisher=publisher) + remote.start() + remote.update_url(url) + + from . import browser_engine + slot_lease = browser_engine._retain_browser_slot(config) + raw_browser = launch_with_driver_cleanup( + launch, + headless=False, + humanize=True, + args=["--disable-features=CrossOriginOpenerPolicy"], + ) + browser = browser_engine._LeasedBrowser(raw_browser, slot_lease) + slot_lease = None context = browser.new_context() page = context.new_page() @@ -288,13 +444,6 @@ def open_login_browser( remote.update_url(current_url) except Exception: log.info(" [browser] Browser closed by user.") - if remote: - remote.stop() - if not keep_alive: - try: - browser.close() - except Exception: - pass return (False, None, None, None) if keep_alive else False if detect_login and detect_login(context, page): @@ -306,11 +455,9 @@ def open_login_browser( print(f" 登录成功！Cookie 已保存至 {cookie_file}") if auto_import: _import_to_browser(netscape_path, config) - if remote: - remote.stop() if keep_alive: - return True, context, page - browser.close() + ownership_transferred = True + return True, browser, context, page return True url_lower = current_url.lower() @@ -324,27 +471,34 @@ def open_login_browser( print(f" 登录成功！Cookie 已保存至 {cookie_file}") if auto_import: _import_to_browser(netscape_path, config) - if remote: - remote.stop() if keep_alive: - return True, context, page - browser.close() + ownership_transferred = True + return True, browser, context, page return True print(" 登录超时。") - if remote: - remote.stop() - if not keep_alive: - try: - browser.close() - except Exception: - pass return (False, None, None, None) if keep_alive else False except Exception as exc: log.info(f" [browser] Login error: {exc}") print(f" 登录出错: {exc}") return (False, None, None, None) if keep_alive else False + finally: + if remote is not None: + try: + remote.stop() + except Exception: + pass + if browser is not None and not ownership_transferred: + try: + browser.close() + except Exception as exc: + log.info( + " [browser] Login browser close failed; retaining global slot: " + f"{exc}" + ) + elif slot_lease is not None: + slot_lease.close() def webvpn_login(config: dict[str, Any]) -> bool: diff --git a/src/scansci_pdf/cache.py b/src/scansci_pdf/cache.py index 9a5cf50..32244e7 100644 --- a/src/scansci_pdf/cache.py +++ b/src/scansci_pdf/cache.py @@ -5,6 +5,7 @@ import hashlib import json import time +import uuid from pathlib import Path from typing import Any @@ -35,7 +36,7 @@ def cache_set(identifier: str, result: dict[str, Any], config: dict[str, Any]) - cache_dir = Path(config["cache_dir"]) cache_dir.mkdir(parents=True, exist_ok=True) target = cache_path(identifier, config) - tmp = target.with_suffix(".tmp") + tmp = target.with_name(f".{target.name}.{uuid.uuid4().hex}.tmp") try: with tmp.open("w", encoding="utf-8") as fh: json.dump(result, fh, indent=2, ensure_ascii=False) diff --git a/src/scansci_pdf/cli.py b/src/scansci_pdf/cli.py index 0737858..e769897 100644 --- a/src/scansci_pdf/cli.py +++ b/src/scansci_pdf/cli.py @@ -367,13 +367,15 @@ def login( _setup_logging(verbose) config = load_config() fetcher = PaperFetcher(config) - - console.print("[bold]Checking institutional access session...[/bold]") - if fetcher.auth.login(force=force): - console.print("[green]Institutional access session is active.[/green]") - else: - console.print("[red]Failed to authenticate institutional access.[/red]") - raise typer.Exit(1) + try: + console.print("[bold]Checking institutional access session...[/bold]") + if fetcher.auth.login(force=force): + console.print("[green]Institutional access session is active.[/green]") + else: + console.print("[red]Failed to authenticate institutional access.[/red]") + raise typer.Exit(1) + finally: + fetcher.close() @app.command() @@ -771,6 +773,25 @@ def session_broker_stop( console.print(f"[green]Stop requested for broker:[/green] {publisher}") +def _fetch_broker_record(downloader, context, record, run_dir): + from .publisher_batch import DownloadResult, _BrowserContextInvalidated + + if context is None: + context = downloader._launch_context() + try: + return downloader.fetch_one(context, record, run_dir), context + except _BrowserContextInvalidated as exc: + return exc.result, None + except Exception as exc: + downloader._close_resource_with_retry(context) + return DownloadResult( + doi=record.doi, + status="failed", + reason=f"{type(exc).__name__}: {exc}", + state="unexpected_error", + ), None + + @app.command("session-broker-run", hidden=True) def session_broker_run( publisher: str = typer.Option(..., "--publisher", "-p"), @@ -835,7 +856,13 @@ def session_broker_run( records = [PaperRecord(**record) for record in job.get("records", [])] results = [] for record in records: - results.append(job_downloader.fetch_one(context, record, primary_dir)) + result, context = _fetch_broker_record( + job_downloader, + context, + record, + primary_dir, + ) + results.append(result) job_downloader._write_results(primary_dir / "summary_partial.json", results) job_downloader._write_results(primary_dir / "summary.json", results) summary = job_downloader._write_complete_artifacts(records, results, run_dir) @@ -858,10 +885,8 @@ def session_broker_run( job_path.unlink(missing_ok=True) time.sleep(2) finally: - try: - context.close() - except Exception: - pass + if context is not None: + downloader._close_resource_with_retry(context) @app.command("session-doctor") diff --git a/src/scansci_pdf/cloakbrowser_compat.py b/src/scansci_pdf/cloakbrowser_compat.py index e7f3cae..0d16b12 100644 --- a/src/scansci_pdf/cloakbrowser_compat.py +++ b/src/scansci_pdf/cloakbrowser_compat.py @@ -4,12 +4,74 @@ import os import platform +import threading from pathlib import Path -from typing import Any +from typing import Any, Callable, TypeVar _CLOAKBROWSER_CACHE_ENV = "CLOAKBROWSER_CACHE_DIR" _SCANSCI_CACHE_ENV = "SCANSCI_PDF_CLOAKBROWSER_CACHE_DIR" _BUILTIN_CACHE_DIR = Path(__file__).resolve().parent / "_browsers" / "cloakbrowser" +_launch_cleanup_lock = threading.RLock() +_LaunchResult = TypeVar("_LaunchResult") + + +class _TrackedPlaywrightContextManager: + """Record Playwright instances started while CloakBrowser is launching.""" + + def __init__(self, manager: Any, started: list[Any]) -> None: + self._manager = manager + self._started = started + + def start(self) -> Any: + playwright = self._manager.start() + self._started.append(playwright) + return playwright + + def __enter__(self) -> Any: + playwright = self._manager.__enter__() + self._started.append(playwright) + return playwright + + def __exit__(self, *args: Any) -> Any: + return self._manager.__exit__(*args) + + def __getattr__(self, name: str) -> Any: + return getattr(self._manager, name) + + +def launch_with_driver_cleanup( + launch_callable: Callable[..., _LaunchResult], + *args: Any, + **kwargs: Any, +) -> _LaunchResult: + """Stop Playwright when CloakBrowser fails after starting its driver. + + CloakBrowser owns the driver after a successful launch and stops it from the + returned browser/context ``close()`` method. CloakBrowser 0.4.10 does not + stop the driver when Chromium launch or post-launch patching raises. + """ + import playwright.sync_api as playwright_sync_api + + with _launch_cleanup_lock: + original_sync_playwright = playwright_sync_api.sync_playwright + started: list[Any] = [] + + def tracked_sync_playwright(*factory_args: Any, **factory_kwargs: Any) -> Any: + manager = original_sync_playwright(*factory_args, **factory_kwargs) + return _TrackedPlaywrightContextManager(manager, started) + + playwright_sync_api.sync_playwright = tracked_sync_playwright + try: + return launch_callable(*args, **kwargs) + except BaseException: + for playwright in reversed(started): + try: + playwright.stop() + except BaseException: + pass + raise + finally: + playwright_sync_api.sync_playwright = original_sync_playwright def configure_builtin_cloakbrowser( diff --git a/src/scansci_pdf/deps.py b/src/scansci_pdf/deps.py index 986ffc0..b759382 100644 --- a/src/scansci_pdf/deps.py +++ b/src/scansci_pdf/deps.py @@ -13,6 +13,7 @@ CORE_DEPS = { "requests": "HTTP client", "bs4": "HTML parsing (beautifulsoup4)", + "pymupdf": "PDF text extraction (PyMuPDF)", "mcp": "MCP protocol", "typer": "CLI framework", "uvicorn": "ASGI server", diff --git a/src/scansci_pdf/extractors/pdf_extractor.py b/src/scansci_pdf/extractors/pdf_extractor.py index 52bed6c..ded2e17 100644 --- a/src/scansci_pdf/extractors/pdf_extractor.py +++ b/src/scansci_pdf/extractors/pdf_extractor.py @@ -4,7 +4,10 @@ import re from pathlib import Path -import pymupdf +try: + import pymupdf +except ImportError: # Allow reduced runtime images to import the package. + pymupdf = None logger = logging.getLogger(__name__) @@ -29,19 +32,27 @@ def extract_text(pdf_path: str | Path) -> str: logger.error("PDF file not found: %s", pdf_path) return "" + if pymupdf is None: + logger.error("PyMuPDF is required for PDF text extraction") + return "" + + doc = None try: doc = pymupdf.open(str(pdf_path)) + text_parts = [] + for page in doc: + text = page.get_text("text") + if text: + text_parts.append(text) except Exception as e: - logger.error("Failed to open PDF %s: %s", pdf_path, e) + logger.error("Failed to extract text from PDF %s: %s", pdf_path, e) return "" - - text_parts = [] - for page in doc: - text = page.get_text("text") - if text: - text_parts.append(text) - - doc.close() + finally: + if doc is not None: + try: + doc.close() + except Exception: + pass full_text = "\n\n".join(text_parts) # Clean up common PDF artifacts @@ -101,19 +112,27 @@ def extract_from_bytes(pdf_bytes: bytes) -> str: Returns: Extracted text content. """ + if pymupdf is None: + logger.error("PyMuPDF is required for PDF text extraction") + return "" + + doc = None try: doc = pymupdf.open(stream=pdf_bytes, filetype="pdf") + text_parts = [] + for page in doc: + text = page.get_text("text") + if text: + text_parts.append(text) except Exception as e: - logger.error("Failed to open PDF from bytes: %s", e) + logger.error("Failed to extract text from PDF bytes: %s", e) return "" - - text_parts = [] - for page in doc: - text = page.get_text("text") - if text: - text_parts.append(text) - - doc.close() + finally: + if doc is not None: + try: + doc.close() + except Exception: + pass return _clean_text("\n\n".join(text_parts)) diff --git a/src/scansci_pdf/fetcher.py b/src/scansci_pdf/fetcher.py index 05c81b9..4798f7e 100644 --- a/src/scansci_pdf/fetcher.py +++ b/src/scansci_pdf/fetcher.py @@ -8,6 +8,7 @@ import re import time from pathlib import Path +from typing import Any from urllib.parse import urlparse import requests @@ -24,7 +25,7 @@ prepare_cloakbrowser_runtime() import cloakbrowser # noqa: F401 _HAS_CLOAKBROWSER = True -except ImportError: +except Exception: _HAS_CLOAKBROWSER = False logger = logging.getLogger(__name__) @@ -135,6 +136,16 @@ def _apply_pdf_bytes(paper: Paper, pdf_bytes: bytes, doi: str, source: str, paper.source = source +def _close_response(response) -> None: + """Release a consumed HTTP response without masking the fetch outcome.""" + if response is None: + return + try: + response.close() + except Exception: + pass + + def _wait_for_challenge(page, max_tries: int = 6) -> None: """Wait for Cloudflare/bot-detection challenges to clear.""" for i in range(max_tries): @@ -384,6 +395,7 @@ def _try_publisher_pdf(self, doi: str, resolved_url: str, paper: Paper) -> Paper return None self._rate_limit() + resp = None try: resp = self.auth.fetch(pdf_url) resp.raise_for_status() @@ -399,6 +411,8 @@ def _try_publisher_pdf(self, doi: str, resolved_url: str, paper: Paper) -> Paper ct, len(resp.content)) except requests.RequestException as e: logger.warning("Failed to fetch publisher PDF: %s", e) + finally: + _close_response(resp) return None @@ -430,15 +444,35 @@ def _try_browser_pdf_download(self, doi: str, resolved_url: str, paper: Paper) - def _browser_pdf_download(self, context, article_url: str, doi: str, paper: Paper) -> Paper | None: page = context.new_page() + cleanup_callbacks = [] try: - return self._browser_pdf_do(page, article_url, doi, paper) + return self._browser_pdf_do( + page, + article_url, + doi, + paper, + cleanup_callbacks=cleanup_callbacks, + ) finally: + for cleanup in reversed(cleanup_callbacks): + try: + cleanup() + except Exception: + pass try: page.close() except Exception: pass - def _browser_pdf_do(self, page, article_url: str, doi: str, paper: Paper) -> Paper | None: + def _browser_pdf_do( + self, + page, + article_url: str, + doi: str, + paper: Paper, + *, + cleanup_callbacks: list[Any] | None = None, + ) -> Paper | None: try: proxied_article = self.auth.convert_url(article_url) except Exception: @@ -527,6 +561,21 @@ def _on_response(response): pass page.on("response", _on_response) + capture_active = True + + def _cleanup_response_capture(): + nonlocal capture_active + if capture_active: + try: + page.remove_listener("response", _on_response) + except Exception: + pass + finally: + capture_active = False + captured_pdf["bytes"] = None + + if cleanup_callbacks is not None: + cleanup_callbacks.append(_cleanup_response_capture) pdf_paths = self._build_browser_pdf_paths(doi, current_url) all_urls = [pdf_link] @@ -583,9 +632,8 @@ def _on_response(response): logger.warning("Anti-bot detected on PDF page, trying next URL") continue - page.remove_listener("response", _on_response) - pdf_bytes = captured_pdf["bytes"] + _cleanup_response_capture() if pdf_bytes and pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: _apply_pdf_bytes(paper, pdf_bytes, doi, "browser", self._save_pdf) logger.info("Browser PDF downloaded via response capture (%d bytes)", len(pdf_bytes)) @@ -623,6 +671,8 @@ def _try_carsi_pdf(self, doi: str, resolved_url: str, paper: Paper) -> Paper | N logger.info("Trying CARSI publisher PDF: %s", pdf_url) self._rate_limit() + carsi = None + resp = None try: carsi = CARSIClient(self.config) resp = carsi.fetch(pdf_url) @@ -635,6 +685,13 @@ def _try_carsi_pdf(self, doi: str, resolved_url: str, paper: Paper) -> Paper | N return result except Exception as e: logger.warning("CARSI PDF failed: %s", e) + finally: + _close_response(resp) + if carsi is not None: + try: + carsi.close() + except Exception as e: + logger.warning("Failed to close CARSI PDF client: %s", e) return None def _try_carsi_html(self, url: str, paper: Paper) -> Paper | None: @@ -642,6 +699,8 @@ def _try_carsi_html(self, url: str, paper: Paper) -> Paper | None: logger.info("Trying CARSI HTML: %s", url) self._rate_limit() + carsi = None + resp = None try: carsi = CARSIClient(self.config) resp = carsi.fetch(url) @@ -660,6 +719,13 @@ def _try_carsi_html(self, url: str, paper: Paper) -> Paper | None: return result except Exception as e: logger.warning("CARSI HTML failed: %s", e) + finally: + _close_response(resp) + if carsi is not None: + try: + carsi.close() + except Exception as e: + logger.warning("Failed to close CARSI HTML client: %s", e) return None def _fetch_via_webvpn(self, url: str, paper: Paper) -> Paper: @@ -669,37 +735,44 @@ def _fetch_via_webvpn(self, url: str, paper: Paper) -> Paper: paper.source = "institutional" + resp = None try: - resp = self.auth.fetch(url) - resp.raise_for_status() - except requests.RequestException as e: - logger.error("Failed to fetch via institutional access: %s", e) - return paper + try: + resp = self.auth.fetch(url) + resp.raise_for_status() + except requests.RequestException as e: + logger.error("Failed to fetch via institutional access: %s", e) + return paper - if "pdf" in resp.headers.get("content-type", "").lower(): - _apply_pdf_bytes(paper, resp.content, paper.doi or "unknown", "institutional", - self._save_pdf) - return paper + if "pdf" in resp.headers.get("content-type", "").lower(): + _apply_pdf_bytes(paper, resp.content, paper.doi or "unknown", "institutional", + self._save_pdf) + return paper - self._apply_extracted(paper, html_extractor.extract(resp.text, resp.url)) + self._apply_extracted(paper, html_extractor.extract(resp.text, resp.url)) - pdf_url = self._find_pdf_link(resp.text, resp.url) - if pdf_url: - logger.info("Found PDF link in HTML, downloading: %s", pdf_url) - self._rate_limit() - try: - pdf_resp = self.auth.fetch(pdf_url) - pdf_resp.raise_for_status() - ct = pdf_resp.headers.get("content-type", "").lower() - if "pdf" in ct and len(pdf_resp.content) > 10000: - pdf_path = self._save_pdf(paper.doi or "unknown", pdf_resp.content) - paper.pdf_path = str(pdf_path) if pdf_path else "" - if not _is_good_result(paper): - paper.full_text = pdf_extractor.extract_from_bytes(pdf_resp.content) - except requests.RequestException as e: - logger.warning("Failed to download PDF: %s", e) + pdf_url = self._find_pdf_link(resp.text, resp.url) + if pdf_url: + logger.info("Found PDF link in HTML, downloading: %s", pdf_url) + self._rate_limit() + pdf_resp = None + try: + pdf_resp = self.auth.fetch(pdf_url) + pdf_resp.raise_for_status() + ct = pdf_resp.headers.get("content-type", "").lower() + if "pdf" in ct and len(pdf_resp.content) > 10000: + pdf_path = self._save_pdf(paper.doi or "unknown", pdf_resp.content) + paper.pdf_path = str(pdf_path) if pdf_path else "" + if not _is_good_result(paper): + paper.full_text = pdf_extractor.extract_from_bytes(pdf_resp.content) + except requests.RequestException as e: + logger.warning("Failed to download PDF: %s", e) + finally: + _close_response(pdf_resp) - return paper + return paper + finally: + _close_response(resp) def _try_elsevier_api(self, doi: str, paper: Paper) -> Paper | None: from .sources import elsevier_api @@ -852,6 +925,7 @@ def _parse_url(self, identifier: str) -> str | None: return identifier if identifier.startswith("http") else None def _resolve_doi(self, doi: str) -> str | None: + resp = None try: resp = request_with_retry( "GET", @@ -861,12 +935,13 @@ def _resolve_doi(self, doi: str) -> str | None: headers={"User-Agent": "scansci-pdf/1.5"}, stream=True, ) - resp.close() if resp.url and resp.url != f"https://doi.org/{doi}": logger.info("Resolved DOI %s → %s (status=%d)", doi, resp.url, resp.status_code) return resp.url except requests.RequestException as e: logger.warning("Failed to resolve DOI %s: %s", doi, e) + finally: + _close_response(resp) return None def _rate_limit(self): @@ -930,5 +1005,7 @@ def clear_cache(self): logger.info("Cache cleared.") def close(self): - if self._auth: - self._auth.close() + auth = self._auth + self._auth = None + if auth is not None: + auth.close() diff --git a/src/scansci_pdf/main.py b/src/scansci_pdf/main.py index b9353a8..ad51ae1 100644 --- a/src/scansci_pdf/main.py +++ b/src/scansci_pdf/main.py @@ -19,7 +19,7 @@ class ServerMode(str, Enum): @app.command("run") def run_server( mode: ServerMode = typer.Option(ServerMode.STDIO, help="Transport mode"), - host: str = typer.Option("0.0.0.0", help="HTTP host"), + host: str = typer.Option("0.0.0.0", help="HTTP bind host"), port: int = typer.Option(8000, help="HTTP port"), ) -> None: """Start the ScanSci PDF server.""" @@ -350,23 +350,25 @@ def fetch_paper_cmd( Cascade: cache → OA → Elsevier API → DOI resolve → CARSI → publisher → browser → gateway. """ - from .institutional.config_adapter import ConfigAdapter - from .institutional.fetcher import PaperFetcher + from .config import load_config + from .fetcher import PaperFetcher - config = ConfigAdapter.load() - config._config["output_dir"] = output + config = load_config() + if output: + config["output_dir"] = output fetcher = PaperFetcher(config) - result = fetcher.fetch_with_result(identifier, use_cache=not no_cache) - - if format == "json": - print(result.to_json()) - elif format == "text": - print(result.to_text()) - else: - print(result.to_markdown(include_pdf_path=True)) + try: + result = fetcher.fetch_with_result(identifier, use_cache=not no_cache) - fetcher.close() + if format == "json": + print(result.to_json()) + elif format == "text": + print(result.to_text()) + else: + print(result.to_markdown(include_pdf_path=True)) + finally: + fetcher.close() @app.command("batch") @@ -378,7 +380,13 @@ def batch_fetch_cmd( ) -> None: """Batch fetch papers. Default: institutional cascade. Use --scihub for grey-source racing.""" import json as _json - from .config import load_config as _load_config + from .config import load_config + from .fetcher import PaperFetcher + + # Typer injects a bool at the CLI boundary. Direct Python callers see the + # OptionInfo default object, which must not be treated as truthy opt-in. + if not isinstance(scihub, bool): + scihub = False dois = [ line.strip() for line in Path(input_file).read_text(encoding="utf-8").splitlines() @@ -389,7 +397,7 @@ def batch_fetch_cmd( return # Auto-detect: if download_strategy is grey/scihub-oriented, switch to racing engine - _cfg = _load_config() + _cfg = load_config() _strategy = _cfg.get("download_strategy", "fastest") _auto_scihub = scihub or _strategy in ("scihub_only", "grey_only", "scihub_first") if not scihub and _auto_scihub: @@ -408,35 +416,36 @@ def batch_fetch_cmd( return # Default: institutional cascade (PaperFetcher) - from .institutional.config_adapter import ConfigAdapter - from .institutional.fetcher import PaperFetcher - - config = ConfigAdapter.load() - config._config["output_dir"] = output + config = load_config() + if output: + config["output_dir"] = output fetcher = PaperFetcher(config) results = [] - for i, doi in enumerate(dois, 1): - print(f" [{i}/{len(dois)}] {doi}") - try: - result = fetcher.fetch_with_result(doi) - result_dict = result.to_dict() - # Verify file actually exists on disk for success status - if result_dict.get("status") == "success" or result_dict.get("success"): - pdf_path = result_dict.get("file") or result_dict.get("pdf_path", "") - if pdf_path and not Path(pdf_path).exists(): - result_dict["status"] = "error" - result_dict["error"] = "PDF file not found on disk (may have been saved elsewhere)" - results.append(result_dict) - status = result.status - quality = result.quality - print(f" → {status} ({quality})") - except Exception as e: - results.append({"doi": doi, "error": str(e)}) - print(f" → error: {e}") - - fetcher.close() + try: + for i, doi in enumerate(dois, 1): + print(f" [{i}/{len(dois)}] {doi}") + try: + result = fetcher.fetch_with_result(doi) + result_dict = result.to_dict() + if result_dict.get("status") == "success" or result_dict.get("success"): + pdf_path = result_dict.get("file") or result_dict.get("pdf_path", "") + if pdf_path and not Path(pdf_path).exists(): + result_dict["status"] = "error" + result_dict["success"] = False + result_dict["error"] = ( + "PDF file not found on disk (may have been saved elsewhere)" + ) + results.append(result_dict) + status = result.status + quality = result.quality + print(f" → {status} ({quality})") + except Exception as e: + results.append({"doi": doi, "error": str(e)}) + print(f" → error: {e}") + finally: + fetcher.close() if format == "json": out_path = Path(output) / "batch_results.json" @@ -497,11 +506,11 @@ def elsevier_setup( @app.command("session-doctor") def session_doctor() -> None: """Diagnose browser profile sessions and cookie health.""" - from .institutional.config_adapter import ConfigAdapter - from .institutional.profile_health import candidate_profile_dirs, inspect_browser_profile + from .config import load_config + from .profile_health import candidate_profile_dirs, inspect_browser_profile - config = ConfigAdapter.load() - profiles = candidate_profile_dirs(config.chrome_profile_dir) + config = load_config() + profiles = candidate_profile_dirs(config) domains = [ "sciencedirect.com", "springer.com", "nature.com", "wiley.com", @@ -538,14 +547,15 @@ def federated_login( config = load_config() client = CARSIClient(config) - - success = client.login(publisher, force=force) - if success: - print(f" Login successful for {publisher}.") - else: - print(f" Login failed for {publisher}.") - raise typer.Exit(1) - client.close() + try: + success = client.login(publisher, force=force) + if success: + print(f" Login successful for {publisher}.") + else: + print(f" Login failed for {publisher}.") + raise typer.Exit(1) + finally: + client.close() @app.command("publisher-batch") @@ -557,6 +567,8 @@ def publisher_batch_cmd( ) -> None: """Batch download papers via publisher-specific workflows.""" from .config import load_config + from .publisher_batch import PaperRecord, PublisherBatchDownloader + from .publisher_profiles import get_publisher_profile, infer_publisher_profile dois = [ line.strip() for line in Path(input_file).read_text(encoding="utf-8").splitlines() @@ -569,16 +581,49 @@ def publisher_batch_cmd( config = load_config() config["output_dir"] = output + records = [PaperRecord(doi=doi) for doi in dois] + if publisher: + try: + profile = get_publisher_profile(publisher) + except ValueError as exc: + print(f" Error: {exc}") + raise typer.Exit(1) from exc + else: + inferred = [infer_publisher_profile(record.doi) for record in records] + profile_names = { + candidate.name for candidate in inferred if candidate is not None + } + if any(candidate is None for candidate in inferred) or len(profile_names) != 1: + print(" Error: could not infer one publisher for all DOIs; use --publisher.") + raise typer.Exit(1) + profile = inferred[0] + + try: + concurrency = max(1, int(max_workers)) + except (TypeError, ValueError) as exc: + print(f" Error: invalid --max-workers value: {max_workers}") + raise typer.Exit(1) from exc + print(f" Batch: {len(dois)} DOIs") - print(f" Publisher: {publisher or 'auto-detect'}") + print(f" Publisher: {profile.name}") print() - # Use the existing publisher batch infrastructure - from .institutional.publisher_batch import PublisherBatchDownloader - downloader = PublisherBatchDownloader(config) - results = downloader.run(dois, publisher=publisher) + run_dir = Path(output or config.get("output_dir", ".")) + institution_query = str( + config.get("carsi_idp_name") or config.get("instsci_school") or "" + ) + downloader = PublisherBatchDownloader( + config, + profile=profile, + institution_query=institution_query, + ) + summary = downloader.run_records( + records, + run_dir, + concurrency=concurrency, + ) - success = sum(1 for r in results if r.get("success")) + success = int(summary.get("success", 0)) print(f"\n Results: {success}/{len(dois)} downloaded") diff --git a/src/scansci_pdf/pdf_utils.py b/src/scansci_pdf/pdf_utils.py index add7b47..f43d9f7 100644 --- a/src/scansci_pdf/pdf_utils.py +++ b/src/scansci_pdf/pdf_utils.py @@ -2,8 +2,11 @@ from __future__ import annotations +import contextlib import re +import threading import urllib.parse +import uuid from pathlib import Path from typing import Any @@ -52,6 +55,8 @@ def is_plausible_pdf_url(url: str) -> bool: return True if "format=pdf" in query or "type=pdf" in query: return True + if "pdf=render" in query: + return True if ("hal.science" in host or "archives-ouvertes" in host) and path.endswith("/document"): return True return False @@ -83,7 +88,7 @@ def is_suspicious_pdf(path: Path) -> bool: # Count PDF page objects: look for "/Type /Page" not followed by "s" import re pages = len(re.findall(rb"/Type\s*/Page\b", content)) - if pages <= 1: + if pages == 1: return True return False except OSError: @@ -101,6 +106,147 @@ def suspicious_pdf(identifier: str, file_path: Path, source_label: str) -> dict[ "error_type": "suspicious_pdf", "reason": "PDF appears to be a cover page or preview (too small / too few pages)", } +def _unique_part_path(output_path: Path) -> Path: + """Return a collision-resistant temporary path beside ``output_path``.""" + return output_path.with_name( + f".{output_path.name}.{uuid.uuid4().hex}.part" + ) + + +def _publish_temp_file_atomic( + tmp_path: Path, + output_path: Path, + cancel_event: threading.Event | None = None, +) -> bool: + """Atomically publish ``tmp_path`` and safely retract our own cancelled write.""" + if cancel_event is not None and cancel_event.is_set(): + return False + + try: + tmp_stat = tmp_path.stat() + if cancel_event is not None and cancel_event.is_set(): + return False + tmp_path.replace(output_path) + if cancel_event is None or not cancel_event.is_set(): + return True + + try: + published_stat = output_path.stat() + identity = ("st_dev", "st_ino", "st_size", "st_mtime_ns") + if all( + getattr(published_stat, name, None) + == getattr(tmp_stat, name, None) + for name in identity + ): + output_path.unlink(missing_ok=True) + except OSError: + pass + except Exception: + return False + return False + + +def publish_pdf_file_atomic( + source_path: Path, + output_path: Path, + cancel_event: threading.Event | None = None, +) -> bool: + """Move a completed PDF into place using the cancellation-safe publisher.""" + if source_path == output_path: + return cancel_event is None or not cancel_event.is_set() + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + except OSError: + return False + return _publish_temp_file_atomic(source_path, output_path, cancel_event) + + +def write_pdf_bytes_atomic( + output_path: Path, + content: bytes, + cancel_event: threading.Event | None = None, +) -> bool: + """Publish browser-captured PDF bytes only after a complete temp-file write.""" + if cancel_event is not None and cancel_event.is_set(): + return False + + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + except OSError: + return False + tmp_path = _unique_part_path(output_path) + try: + with tmp_path.open("wb") as fh: + if cancel_event is not None and cancel_event.is_set(): + raise InterruptedError("PDF write cancelled") + fh.write(content) + if cancel_event is not None and cancel_event.is_set(): + raise InterruptedError("PDF write cancelled") + return _publish_temp_file_atomic(tmp_path, output_path, cancel_event) + except Exception: + return False + finally: + with contextlib.suppress(OSError): + tmp_path.unlink(missing_ok=True) + + +def write_pdf_stream_atomic( + output_path: Path, + first_chunk: bytes, + chunks: Any, + cancel_event: threading.Event | None = None, +) -> bool: + """Write a streamed PDF to a unique temporary file and publish atomically.""" + if cancel_event is not None and cancel_event.is_set(): + return False + + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + except OSError: + return False + tmp_path = _unique_part_path(output_path) + try: + with tmp_path.open("wb") as fh: + if cancel_event is not None and cancel_event.is_set(): + raise InterruptedError("PDF write cancelled") + fh.write(first_chunk) + for chunk in chunks: + if cancel_event is not None and cancel_event.is_set(): + raise InterruptedError("PDF write cancelled") + if chunk: + fh.write(chunk) + if cancel_event is not None and cancel_event.is_set(): + raise InterruptedError("PDF write cancelled") + return _publish_temp_file_atomic(tmp_path, output_path, cancel_event) + except Exception: + return False + finally: + with contextlib.suppress(OSError): + tmp_path.unlink(missing_ok=True) + + +def _bind_session_to_response( + response: requests.Response, + session: requests.Session, +) -> requests.Response: + """Make closing an escaping response also close its private session.""" + original_close = response.close + close_lock = threading.Lock() + closed = False + + def close() -> None: + nonlocal closed + with close_lock: + if closed: + return + closed = True + try: + original_close() + finally: + session.close() + + response.close = close # type: ignore[method-assign] + return response def success(identifier: str, file_path: Path, source: str) -> dict[str, Any]: @@ -193,10 +339,15 @@ def download_pdf( require_pdf_like_url: bool = True, use_tor: bool = False, cookies: Any = None, + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: + if cancel_event is not None and cancel_event.is_set(): + return None if require_pdf_like_url and not is_plausible_pdf_url(url): return None + session: requests.Session | None = None + resp: requests.Response | None = None try: if cookies is not None: from .network import request_timeout, proxy_dict, select_proxy_for_url, USER_AGENT @@ -212,35 +363,42 @@ def download_pdf( stream=True, ) else: - resp = fetch(url, config, stream=True, use_tor=use_tor) - if resp.status_code >= 400: + resp = fetch( + url, + config, + stream=True, + use_tor=use_tor, + cancel_event=cancel_event, + ) + if ( + cancel_event is not None and cancel_event.is_set() + ) or resp.status_code >= 400: return None - iterator = resp.iter_content(chunk_size=8192) + iterator = resp.iter_content(chunk_size=65536) first_chunk = next(iterator, b"") if not _response_looks_pdf(resp, first_chunk): return None - output_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = output_path.with_suffix(output_path.suffix + ".part") - try: - with tmp_path.open("wb") as fh: - fh.write(first_chunk) - for chunk in iterator: - if chunk: - fh.write(chunk) - tmp_path.replace(output_path) - except Exception: - tmp_path.unlink(missing_ok=True) - raise + if not write_pdf_stream_atomic( + output_path, + first_chunk, + iterator, + cancel_event, + ): + return None if is_pdf_file(output_path): return success(output_path.stem, output_path, source) - else: - try: - output_path.unlink(missing_ok=True) - except OSError: - pass + try: + output_path.unlink(missing_ok=True) + except OSError: + pass except Exception: return None + finally: + if resp is not None: + resp.close() + if session is not None: + session.close() return None diff --git a/src/scansci_pdf/publisher_batch.py b/src/scansci_pdf/publisher_batch.py index e6bad28..86aac7a 100644 --- a/src/scansci_pdf/publisher_batch.py +++ b/src/scansci_pdf/publisher_batch.py @@ -6,6 +6,7 @@ import csv import inspect import json +import queue import re import shutil import threading @@ -18,7 +19,7 @@ import requests -from .config import load_config, parse_proxy_pool +from .config import DATA_DIR, load_config, parse_proxy_pool from .extractors import pdf_extractor from .publisher_pdf_router import ( build_pdf_candidates, @@ -106,6 +107,14 @@ def ok(self) -> bool: return self.status == "success" +class _BrowserContextInvalidated(RuntimeError): + """Signals that an owned page could not be closed and its context was retired.""" + + def __init__(self, result: DownloadResult, detail: str): + super().__init__(detail) + self.result = result + + def safe_name(doi: str) -> str: return re.sub(r"[^\w\-.]", "_", doi) @@ -162,91 +171,6 @@ def fetch_est_records( return records[:limit] -class _PagePool: - """Thread-safe page pool sharing a single browser context. - - All pages share the same context, which means they share: - - Cookies (session tokens, auth cookies) - - localStorage (publisher auth tokens) - - sessionStorage (temporary session data) - - HTTP connection pool (keep-alive authenticated connections) - - TLS session tickets (established encrypted sessions) - - This is essential: login state is the ENTIRE browser context, not just cookies. - Creating a new context is like switching devices — publishers may reject it. - """ - - def __init__(self, context: Any, max_size: int = 2): - self._context = context - self._max_size = max_size - self._available: list[Any] = [] - self._in_use: set[int] = set() - self._lock = threading.Lock() - self._not_empty = threading.Condition(self._lock) - - def acquire(self) -> Any: - """Acquire a page from the pool. Creates new if pool is empty and under limit.""" - with self._not_empty: - # Try to reuse an available page - while self._available: - page = self._available.pop() - try: - _ = page.url # Test if page is still alive - self._in_use.add(id(page)) - return page - except Exception: - try: - page.close() - except Exception: - pass - - # Create new page if under limit - if len(self._in_use) < self._max_size: - page = self._context.new_page() - self._in_use.add(id(page)) - return page - - # Wait for a page to be returned - while not self._available: - self._not_empty.wait(timeout=30) - if self._available: - page = self._available.pop() - try: - _ = page.url - self._in_use.add(id(page)) - return page - except Exception: - try: - page.close() - except Exception: - pass - # Should not reach here, but just in case - page = self._context.new_page() - self._in_use.add(id(page)) - return page - - def release(self, page: Any) -> None: - """Return a page to the pool for reuse.""" - with self._not_empty: - self._in_use.discard(id(page)) - try: - self._available.append(page) - except Exception: - pass - self._not_empty.notify() - - def close_all(self) -> None: - """Close all pages.""" - with self._lock: - for page in self._available: - try: - page.close() - except Exception: - pass - self._available.clear() - self._in_use.clear() - - class PublisherBatchDownloader: """Deterministic publisher workflow with diagnostic packets for surprises.""" @@ -293,7 +217,15 @@ def run_records( run_path = Path(run_dir) run_path.mkdir(parents=True, exist_ok=True) target = target_verified if target_verified and target_verified > 0 else None - worker_count = min(max(1, int(concurrency or 1)), MAX_BROWSER_CONCURRENCY) + try: + browser_limit = max(1, int(self.config.get("max_browser_workers", 1))) + except (TypeError, ValueError): + browser_limit = 1 + worker_count = min( + max(1, int(concurrency or 1)), + MAX_BROWSER_CONCURRENCY, + browser_limit, + ) if target: worker_count = 1 attempt_cache_path = Path(attempt_cache) if attempt_cache else run_path / "attempts.jsonl" @@ -339,7 +271,7 @@ def run_records( failed_records = [ record - for record, result in zip(records, results) + for record, result in zip(records_to_run, results) if result.status == "failed" and result.reason in RETRYABLE_REASONS ] @@ -376,7 +308,7 @@ def run_records( summary["cached_skipped"] = cached_skipped summary["attempt_cache"] = str(attempt_cache_path) summary["concurrency"] = worker_count - summary["browser_profile_dir"] = str(Path(self.config.get("chrome_profile_dir", ""))) + summary["browser_profile_dir"] = str(self._default_profile_dir()) # Surface auto-stop so callers (CLI/MCP) can tell a halt from a clean # finish. ip_block_count is how many records came back ip_blocked in the # primary pass (the trip trigger). @@ -404,6 +336,9 @@ def _run_once( concurrency: int = 1, ) -> list[DownloadResult]: run_dir.mkdir(parents=True, exist_ok=True) + if not records: + self._write_results(run_dir / "summary.json", []) + return [] worker_count = min(max(1, int(concurrency or 1)), len(records) or 1) if worker_count > 1 and not target_verified: return self._run_once_parallel( @@ -418,10 +353,16 @@ def _run_once( verified_count = 0 consecutive_ip_blocks = 0 - context = self._launch_context() + context: Any | None = self._launch_context() try: - for record in records: - result = self.fetch_one(context, record, run_dir) + for record_index, record in enumerate(records): + context_invalidated = False + try: + result = self.fetch_one(context, record, run_dir) + except _BrowserContextInvalidated as exc: + result = exc.result + context_invalidated = True + context = None results.append(result) if result.ok and result.verified_match: verified_count += 1 @@ -438,11 +379,12 @@ def _run_once( consecutive_ip_blocks = 0 if target_verified and verified_count >= target_verified: break + if context_invalidated: + if record_index + 1 < len(records): + context = self._launch_context() finally: - try: - context.close() - except Exception: - pass + if context is not None: + self._close_resource_with_retry(context) self._write_results(run_dir / "summary.json", results) return results @@ -456,18 +398,10 @@ def _run_once_parallel( attempt_cache_path: Path | None = None, phase: str = "primary", ) -> list[DownloadResult]: - """Run downloads in parallel using a shared browser context + page pool. - - Architecture: - - ONE persistent context (preserves cookies, localStorage, sessionStorage, - HTTP connections, TLS sessions — the complete login state) - - N worker threads, each borrows a page from the pool, downloads, returns it - - All pages share the same context = same identity, no re-login needed - - This is the correct model: login state is NOT just cookies. It includes - localStorage tokens, sessionStorage, connection pools, and TLS sessions. - Exporting cookies to a new context is like "switching devices" — some - publishers reject that. Sharing the context keeps the same device identity. + """Run downloads with one persistent context owned by each worker thread. + + Playwright's sync API is thread-bound. Each worker therefore launches and + closes its own context from a copy of the authenticated profile. """ # ── proxy-pool rotation branch ── proxies = parse_proxy_pool(self.config.get("proxy_pool", "")) @@ -480,8 +414,7 @@ def _run_once_parallel( run_dir.mkdir(parents=True, exist_ok=True) profile_root = run_dir / "worker-profiles" - source_profile = Path(self.config.get("chrome_profile_dir", "")) - profile_dir = self._prepare_worker_profile(source_profile, profile_root / f"{phase}-shared") + source_profile = self._default_profile_dir() results_by_index: dict[int, DownloadResult] = {} results_lock = threading.Lock() @@ -515,34 +448,50 @@ def record_result(index: int, result: DownloadResult) -> None: with count_lock: ip_block_count["n"] = 0 - # Single shared context — login state fully preserved - context = self._launch_context(profile_dir=profile_dir) - page_pool = _PagePool(context, max_size=worker_count) - - def run_worker(items: list[tuple[int, PaperRecord]]) -> None: - for item_index, record in items: - if stop_event.is_set(): - break # IP block tripped — skip the rest of this chunk - page = page_pool.acquire() - try: - record_result(item_index, self.fetch_one(page, record, run_dir)) - finally: - page_pool.release(page) - indexed_records = list(enumerate(records)) chunks = [indexed_records[i::worker_count] for i in range(worker_count)] + worker_inputs = [] + for worker_index, chunk in enumerate(chunks): + if not chunk: + continue + profile_dir = self._prepare_worker_profile( + source_profile, + profile_root / f"{phase}-worker-{worker_index}", + ) + worker_inputs.append((chunk, profile_dir)) + + def run_worker( + items: list[tuple[int, PaperRecord]], + profile_dir: Path, + ) -> None: + context: Any | None = self._launch_context(profile_dir=profile_dir) + try: + for item_offset, (item_index, record) in enumerate(items): + if stop_event.is_set(): + break + context_invalidated = False + try: + result = self.fetch_one(context, record, run_dir) + except _BrowserContextInvalidated as exc: + result = exc.result + context_invalidated = True + record_result(item_index, result) + if context_invalidated: + context = None + if item_offset + 1 < len(items): + context = self._launch_context(profile_dir=profile_dir) + finally: + if context is not None: + self._close_resource_with_retry(context) with ThreadPoolExecutor(max_workers=worker_count) as executor: - futures = [executor.submit(run_worker, chunk) for chunk in chunks if chunk] + futures = [ + executor.submit(run_worker, chunk, profile_dir) + for chunk, profile_dir in worker_inputs + ] for future in as_completed(futures): future.result() - page_pool.close_all() - try: - context.close() - except Exception: - pass - results = [results_by_index[i] for i in range(len(records)) if i in results_by_index] self._write_results(run_dir / "summary.json", results) return results @@ -557,119 +506,147 @@ def _run_once_parallel_rotating( attempt_cache_path: Path | None = None, phase: str = "primary", ) -> list[DownloadResult]: - """Run downloads across multiple per-proxy contexts (IP rotation). - - - Logs in once, exports cookies, injects into each proxy context. - - Round-robin assigns records to proxies. - - Per-proxy IP-block tracking: when a proxy hits the block threshold, - it's excluded; when all proxies are excluded the run auto-stops. - """ + """Rotate proxies without sharing Playwright contexts across threads.""" run_dir.mkdir(parents=True, exist_ok=True) profile_root = run_dir / "worker-profiles" - source_profile = Path(self.config.get("chrome_profile_dir", "")) + source_profile = self._default_profile_dir() - # ── 1. shared cookies (login once, reuse across proxies) ── login_profile = profile_root / f"{phase}-login" cookies = self._login_and_export_cookies(run_dir, login_profile) - # ── 2. per-proxy context + page-pool ── class _ProxySlot: - __slots__ = ("proxy", "context", "pool", "blocked", "ip_blocks") - def __init__(self, proxy: str, context, pool): + __slots__ = ("index", "proxy", "blocked", "ip_blocks") + + def __init__(self, index: int, proxy: str): + self.index = index self.proxy = proxy - self.context = context - self.pool = pool self.blocked = False self.ip_blocks = 0 - slots: list[_ProxySlot] = [] - self._proxy_blocked = [] # populated as proxies are excluded - for i, proxy in enumerate(proxies): - p_dir = self._prepare_worker_profile(source_profile, profile_root / f"{phase}-p{i}") - ctx = self._launch_context(profile_dir=p_dir, proxy=proxy) - if cookies: - self._inject_cookies(ctx, cookies) - pool = _PagePool(ctx, max_size=max(1, max(1, worker_count) // len(proxies) + 1)) - slots.append(_ProxySlot(proxy, ctx, pool)) - - # ── 3. shared orchestrator state ── + slots = [_ProxySlot(index, proxy) for index, proxy in enumerate(proxies)] + active_workers = min(max(1, worker_count), len(slots)) + slot_groups = [slots[index::active_workers] for index in range(active_workers)] + worker_inputs: list[list[tuple[_ProxySlot, Path]]] = [] + for worker_index, worker_slots in enumerate(slot_groups): + worker_inputs.append( + [ + ( + slot, + self._prepare_worker_profile( + source_profile, + profile_root / f"{phase}-worker-{worker_index}-p{slot.index}", + ), + ) + for slot in worker_slots + ] + ) + + seed_count = min(len(slots), len(records)) + work_queue: queue.Queue[tuple[int, PaperRecord]] = queue.Queue() + for indexed_record in enumerate(records[seed_count:], start=seed_count): + work_queue.put(indexed_record) + results_by_index: dict[int, DownloadResult] = {} results_lock = threading.Lock() attempt_lock = threading.Lock() - rr_lock = threading.Lock() - rr_index = 0 - # No global consecutive-block counter here: in rotation mode each proxy - # is tracked independently (slot.ip_blocks). The whole run only stops - # when ALL proxies are excluded — that's the point of having a pool. + state_lock = threading.Lock() stop_event = threading.Event() + self._proxy_blocked = [] - def _pick_slot() -> _ProxySlot | None: - nonlocal rr_index - with rr_lock: - active = [s for s in slots if not s.blocked] - if not active: - return None - slot = active[rr_index % len(active)] - rr_index += 1 - return slot - - def _record(index: int, result: DownloadResult, slot: _ProxySlot | None) -> None: + def record_result(index: int, result: DownloadResult, slot: _ProxySlot) -> None: with results_lock: results_by_index[index] = result - part = [results_by_index[ii] for ii in sorted(results_by_index)] - self._write_results(run_dir / "summary_partial.json", part) + partial = [results_by_index[item] for item in sorted(results_by_index)] + self._write_results(run_dir / "summary_partial.json", partial) with attempt_lock: self._append_attempt(attempt_cache_path, result, phase) - if result.reason == "ip_blocked": - if slot is not None: + with state_lock: + if result.reason == "ip_blocked": slot.ip_blocks += 1 if slot.ip_blocks >= IP_BLOCK_STOP_THRESHOLD and not slot.blocked: slot.blocked = True self._proxy_blocked.append(slot.proxy) - # If every proxy is now blocked, stop the whole run. - if all(s.blocked for s in slots) and not stop_event.is_set(): + if all(item.blocked for item in slots): stop_event.set() self._ip_block_stopped = True - elif result.ok or result.reason: - # A success (or non-block failure) resets this proxy's streak. - if slot is not None: + elif result.ok or result.reason: slot.ip_blocks = 0 - def _worker(chunk: list[tuple[int, PaperRecord]]) -> None: - for item_index, record in chunk: - if stop_event.is_set(): - break - slot = _pick_slot() - if slot is None: # all proxies excluded - break - page = slot.pool.acquire() + def run_worker(owned_slots: list[tuple[_ProxySlot, Path]]) -> None: + contexts: dict[_ProxySlot, Any | None] = { + slot: None for slot, _profile_dir in owned_slots + } + profile_dirs = {slot: profile_dir for slot, profile_dir in owned_slots} + cursor = 0 + + def process_record( + slot: _ProxySlot, + item_index: int, + record: PaperRecord, + ) -> None: + context = contexts[slot] + if context is None: + context = self._launch_context( + profile_dir=profile_dirs[slot], + proxy=slot.proxy, + ) + contexts[slot] = context + if cookies: + self._inject_cookies(context, cookies) + try: - _record(item_index, self.fetch_one(page, record, run_dir), slot) - finally: - slot.pool.release(page) + result = self.fetch_one(context, record, run_dir) + except _BrowserContextInvalidated as exc: + result = exc.result + contexts[slot] = None + record_result(item_index, result, slot) - # ── 4. execution ── - indexed = list(enumerate(records)) - chunks = [indexed[i::worker_count] for i in range(worker_count)] - with ThreadPoolExecutor(max_workers=worker_count) as executor: - futures = [executor.submit(_worker, c) for c in chunks if c] + try: + # Seed each configured proxy before workers compete for the shared + # queue. Without this, a fast worker can drain short batches before + # another owner thread gets scheduled, defeating proxy rotation. + for slot, _profile_dir in owned_slots: + if slot.index >= seed_count: + continue + process_record(slot, slot.index, records[slot.index]) + + while not stop_event.is_set(): + with state_lock: + available = [slot for slot, _profile_dir in owned_slots if not slot.blocked] + if not available: + break + try: + item_index, record = work_queue.get_nowait() + except queue.Empty: + break + + slot = available[cursor % len(available)] + cursor += 1 + process_record(slot, item_index, record) + work_queue.task_done() + finally: + for context in contexts.values(): + if context is not None: + self._close_resource_with_retry(context) + + with ThreadPoolExecutor(max_workers=active_workers) as executor: + futures = [executor.submit(run_worker, worker_slots) for worker_slots in worker_inputs] for future in as_completed(futures): future.result() - # ── 5. cleanup ── - for slot in slots: - slot.pool.close_all() - try: - slot.context.close() - except Exception: - pass - - results = [results_by_index[i] for i in range(len(records)) if i in results_by_index] + results = [results_by_index[index] for index in range(len(records)) if index in results_by_index] self._write_results(run_dir / "summary.json", results) return results def _prepare_worker_profile(self, source: Path, target: Path) -> Path: + source = source.resolve() + target = target.resolve() + if target == source or source in target.parents or target in source.parents: + raise ValueError( + "worker profile source and target must not contain each other: " + f"source={source}, target={target}" + ) if target.exists(): shutil.rmtree(target, ignore_errors=True) target.parent.mkdir(parents=True, exist_ok=True) @@ -760,15 +737,9 @@ def _login_and_export_cookies(self, run_dir: Path, profile_dir: Path) -> list[di print(f" ✅ Login successful! Exported {len(cookies)} cookies for {len(cookies)} workers") return cookies finally: - try: - page.close() - except Exception: - pass + self._close_resource_with_retry(page) finally: - try: - context.close() - except Exception: - pass + self._close_resource_with_retry(context) def _inject_cookies(self, context: Any, cookies: list[dict]) -> None: """Inject previously exported cookies into a context (Playwright add_cookies). @@ -824,7 +795,7 @@ def _looks_like_logged_in(self, page: Any, current_url: str) -> bool: def _launch_context(self, profile_dir: str | Path | None = None, *, proxy: str | None = None): from .browser_engine import get_persistent_context - profile_path = Path(profile_dir) if profile_dir else Path(self.config.get("chrome_profile_dir", "")) + profile_path = Path(profile_dir) if profile_dir else self._default_profile_dir() profile_path.mkdir(parents=True, exist_ok=True) # Overlay a per-launch proxy without mutating the shared config object. config = dict(self.config) @@ -837,22 +808,48 @@ def _launch_context(self, profile_dir: str | Path | None = None, *, proxy: str | self._context_proxy[id(ctx)] = proxy return ctx + def _default_profile_dir(self) -> Path: + configured = self.config.get("chrome_profile_dir") + if isinstance(configured, str): + configured = configured.strip() + if configured: + return Path(configured) + elif configured is not None: + try: + return Path(configured) + except (TypeError, ValueError): + pass + return DATA_DIR / "browser_profiles" / safe_name(self.profile.name.lower()) + def fetch_one(self, context_or_page: Any, record: PaperRecord, run_dir: Path) -> DownloadResult: - # Support both context (creates new page) and pre-created page (from page pool) - if hasattr(context_or_page, "new_page") and hasattr(context_or_page, "cookies"): - # It's a context — create a new page - page = context_or_page.new_page() - _owns_page = True - else: - # It's already a page (from page pool) - page = context_or_page - _owns_page = False result = DownloadResult( doi=record.doi, status="failed", state="started", article_url=self.profile.article_url(record.doi), ) + page = context_or_page + _owns_page = False + + # Support both context (creates new page) and pre-created page (from page pool) + if hasattr(context_or_page, "new_page") and hasattr(context_or_page, "cookies"): + # It's a context — create a new page + _owns_page = True + try: + page = context_or_page.new_page() + except Exception as exc: + result.reason = f"{type(exc).__name__}: {exc}" + result.state = "unexpected_error" + detail = f"page creation failed: {exc}" + context_close_error = self._close_resource_with_retry(context_or_page) + if context_close_error is not None: + detail += ( + "; context close failed after retry: " + f"{context_close_error}" + ) + self._event(result, "browser_context_invalidated", detail) + raise _BrowserContextInvalidated(result, detail) from exc + try: self._event(result, "article_open", result.article_url) if not self._ensure_login(page, result): @@ -985,12 +982,34 @@ def fetch_one(self, context_or_page: Any, record: PaperRecord, run_dir: Path) -> self._write_diagnostic(page, result, run_dir) return result finally: - self._hold_after_run(page, result) - if _owns_page: - try: - page.close() - except Exception: - pass + try: + self._hold_after_run(page, result) + finally: + if _owns_page: + page_close_error = self._close_resource_with_retry(page) + if page_close_error is not None: + context_close_error = self._close_resource_with_retry( + context_or_page + ) + detail = f"page close failed after retry: {page_close_error}" + if context_close_error is not None: + detail += ( + "; context close failed after retry: " + f"{context_close_error}" + ) + self._event(result, "browser_context_invalidated", detail) + raise _BrowserContextInvalidated(result, detail) from page_close_error + + @staticmethod + def _close_resource_with_retry(resource: Any) -> Exception | None: + last_error = None + for _attempt in range(2): + try: + resource.close() + return None + except Exception as exc: + last_error = exc + return last_error def _hold_after_login(self, page: Any, result: DownloadResult) -> None: if not self.post_login_hold_sec: @@ -2764,7 +2783,7 @@ def _write_diagnostic(self, page: Any, result: DownloadResult, run_dir: Path) -> packet = { **asdict(result), "publisher": self.profile.name, - "browser_profile_dir": str(Path(self.config.get("chrome_profile_dir", ""))), + "browser_profile_dir": str(self._default_profile_dir()), "body_excerpt": self._body_text(page, 2_000), "created_at": datetime.now().isoformat(timespec="seconds"), } diff --git a/src/scansci_pdf/publisher_pdf_router.py b/src/scansci_pdf/publisher_pdf_router.py index d3dd79f..19ba46a 100644 --- a/src/scansci_pdf/publisher_pdf_router.py +++ b/src/scansci_pdf/publisher_pdf_router.py @@ -51,6 +51,7 @@ "2077-0375": "membranes", "2227-7390": "math", "2304-8158": "foods", + "2673-3986": "epidemiologia", } MDPI_JOURNAL_CODE_ISSNS = { journal_code: issn diff --git a/src/scansci_pdf/publisher_strategies.py b/src/scansci_pdf/publisher_strategies.py index cd0378c..24d19a6 100644 --- a/src/scansci_pdf/publisher_strategies.py +++ b/src/scansci_pdf/publisher_strategies.py @@ -15,14 +15,18 @@ from pathlib import Path from typing import Any +from .cloakbrowser_compat import launch_with_driver_cleanup, prepare_cloakbrowser_runtime + try: + prepare_cloakbrowser_runtime() from cloakbrowser import launch, launch_persistent_context _HAS_CLOAKBROWSER = True -except ImportError: +except Exception: launch = None # type: ignore[assignment] launch_persistent_context = None # type: ignore[assignment] _HAS_CLOAKBROWSER = False from .log import get_logger +from .pdf_utils import write_pdf_bytes_atomic log = get_logger() @@ -35,6 +39,37 @@ _visible_browser_active = False +class _BrowserCancelled(RuntimeError): + pass + + +def _cancelled(cancel_event: threading.Event | None) -> bool: + return cancel_event is not None and cancel_event.is_set() + + +def _wait_or_cancel( + cancel_event: threading.Event | None, + timeout: float, +) -> bool: + if cancel_event is None: + time.sleep(max(0.0, timeout)) + return False + return cancel_event.wait(max(0.0, timeout)) + + +def _close_browser_resource(resource: Any) -> bool: + """Best-effort close on the resource's owner thread, with one retry.""" + if resource is None: + return True + for _attempt in range(2): + try: + resource.close() + return True + except Exception: + pass + return False + + def get_last_error() -> tuple[str, str]: """Return (error_type, action) from last _browser_download call.""" return _last_error_type, _last_error_action @@ -78,42 +113,86 @@ def _restore_cookies_to_context(context: Any, config: dict[str, Any]) -> None: @contextlib.contextmanager -def _visible_browser(config: dict[str, Any], publisher: str, *, viewport: dict | None = None): +def _visible_browser( + config: dict[str, Any], + publisher: str, + *, + viewport: dict | None = None, + cancel_event: threading.Event | None = None, +): """Open visible CloakBrowser with persistent profile. Falls back to ephemeral.""" if not _HAS_CLOAKBROWSER: raise RuntimeError("cloakbrowser not installed. Run: pip install cloakbrowser") + from . import browser_engine + from .browser_engine import _build_browser_args profile_dir = _get_profile_dir(config, publisher) browser = None + ctx = None + page = None + slot_lease = None + args = _build_browser_args(config) + if _cancelled(cancel_event): + raise _BrowserCancelled("browser operation cancelled before launch") try: - ctx = launch_persistent_context( + slot_lease = browser_engine._retain_browser_slot(config, cancel_event) + raw_ctx = launch_with_driver_cleanup( + launch_persistent_context, str(profile_dir), headless=False, humanize=True, - args=["--disable-features=CrossOriginOpenerPolicy"], + args=args, ) + ctx = browser_engine._LeasedPersistentContext(raw_ctx, slot_lease) + slot_lease = None page = ctx.new_page() log.info(f" [{publisher}] persistent browser profile: {profile_dir}") # Ensure cookies are loaded from saved file _restore_cookies_to_context(ctx, config) except Exception as _e: + if ctx is not None: + if not _close_browser_resource(ctx): + raise RuntimeError( + f"{publisher} persistent context failed during setup and " + "could not be closed; replacement browser refused" + ) from _e + ctx = None + elif slot_lease is not None: + slot_lease.close() + slot_lease = None + if _cancelled(cancel_event): + raise _BrowserCancelled("browser operation cancelled during launch") log.info(f" [{publisher}] persistent context unavailable ({_e}), using ephemeral") _vp = viewport or {"width": 1440, "height": 900} - browser = launch(headless=False, humanize=True, - args=["--disable-features=CrossOriginOpenerPolicy"]) - ctx = browser.new_context(viewport=_vp) - _restore_cookies_to_context(ctx, config) - page = ctx.new_page() + try: + slot_lease = browser_engine._retain_browser_slot(config, cancel_event) + raw_browser = launch_with_driver_cleanup( + launch, + headless=False, + humanize=True, + args=args, + ) + browser = browser_engine._LeasedBrowser(raw_browser, slot_lease) + slot_lease = None + ctx = browser.new_context(viewport=_vp) + _restore_cookies_to_context(ctx, config) + page = ctx.new_page() + except Exception: + if ctx is not None: + _close_browser_resource(ctx) + if browser is not None: + _close_browser_resource(browser) + elif slot_lease is not None: + slot_lease.close() + raise try: + if _cancelled(cancel_event): + raise _BrowserCancelled("browser operation cancelled after launch") yield ctx, page finally: - try: - if browser: - browser.close() - else: - ctx.close() - except Exception: - pass + _close_browser_resource(page) + _close_browser_resource(ctx) + _close_browser_resource(browser) def _save_all_cookie_formats( @@ -248,7 +327,12 @@ def _inject_cookies_to_tab(tab_id: str, config: dict[str, Any], publisher: str) log.info(f" [{publisher}] imported {total} cookies into browser session") -def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str) -> bool: +def _try_institutional_login( + tab_id: str, + config: dict[str, Any], + publisher: str, + cancel_event: threading.Event | None = None, +) -> bool: """Try institutional login (OpenAthens/CARSI) when paywall detected. Works within the existing browser-engine tab: @@ -261,6 +345,8 @@ def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str """ from .browser_engine import evaluate_js, navigate_tab + if _cancelled(cancel_event): + return False idp_name = config.get("carsi_idp_name", "") if not idp_name: log.info(f" [{publisher}] no carsi_idp_name configured, skipping institutional login") @@ -287,12 +373,13 @@ def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str })() """, config) - if not sso_clicked: + if _cancelled(cancel_event) or not sso_clicked: log.info(f" [{publisher}] no SSO/institutional login link found on page") return False log.info(f" [{publisher}] clicked institutional login: {str(sso_clicked)[:60]}") - time.sleep(8) + if _wait_or_cancel(cancel_event, 8): + return False # Step 2: Look for institution search box and search search_done = evaluate_js(tab_id, f""" @@ -313,10 +400,13 @@ def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str return null; }})() """, config) + if _cancelled(cancel_event): + return False if search_done: log.info(f" [{publisher}] searched for '{idp_en}' via {search_done}") - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return False # Click matching institution result clicked = evaluate_js(tab_id, f""" @@ -333,10 +423,13 @@ def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str return null; }})() """, config) + if _cancelled(cancel_event): + return False if clicked: log.info(f" [{publisher}] selected institution: {clicked}") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return False else: log.info(f" [{publisher}] no matching institution found for '{idp_en}'") return False @@ -350,6 +443,8 @@ def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str current_url = evaluate_js(tab_id, "window.location.href", config) or "" current_title = evaluate_js(tab_id, "document.title", config) or "" + if _cancelled(cancel_event): + return False needs_login = any(x in current_url.lower() for x in _ak) or any(x in current_title for x in _at) @@ -357,7 +452,15 @@ def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str # browser-engine is headless — user can't see the tab # Open a visible browser window for the CAS login log.info(f" [{publisher}] CAS login required — opening visible browser...") - return _visible_institutional_login(current_url, config, publisher, idp_name, idp_en, tab_id) + return _visible_institutional_login( + current_url, + config, + publisher, + idp_name, + idp_en, + tab_id, + cancel_event=cancel_event, + ) # If we ended up back on the article page, login might have succeeded via cookies log.info(f" [{publisher}] institutional login flow completed") @@ -367,11 +470,19 @@ def _try_institutional_login(tab_id: str, config: dict[str, Any], publisher: str def _visible_institutional_login( cas_url: str, config: dict[str, Any], publisher: str, idp_name: str, idp_en: str, headless_tab_id: str, + cancel_event: threading.Event | None = None, ) -> bool: """Open a visible browser for CAS login, then inject cookies back.""" from .browser_engine import evaluate_js, navigate_tab, import_cookies - with _visible_browser(config, publisher, viewport=None) as (context, page): + if _cancelled(cancel_event): + return False + with _visible_browser( + config, + publisher, + viewport=None, + cancel_event=cancel_event, + ) as (context, page): # Start from article page to get Cloudflare clearance, then do SSO flow article_url = evaluate_js(headless_tab_id, "window.location.href", config) or cas_url @@ -383,11 +494,14 @@ def _visible_institutional_login( article_page = f"https://{publisher_host}/doi/{doi_str}" if doi_str else article_url log.info(f" [{publisher}] visible browser: loading article page first...") + if _cancelled(cancel_event): + return False try: page.goto(article_page, wait_until="domcontentloaded", timeout=60000) - time.sleep(5) except Exception: pass + if _wait_or_cancel(cancel_event, 5): + return False title = page.title() url = page.url @@ -405,7 +519,8 @@ def _visible_institutional_login( if (a) a.click(); })() """) - time.sleep(8) + if _wait_or_cancel(cancel_event, 8): + return False title = page.title() url = page.url @@ -415,7 +530,8 @@ def _visible_institutional_login( si = page.query_selector('#searchInstitution') if si: si.fill(idp_en) - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return False page.evaluate(f""" (name) => {{ const items = document.querySelectorAll('[class*="result"], [class*="suggestion"], li, a, button'); @@ -428,14 +544,16 @@ def _visible_institutional_login( return false; }} """, idp_en) - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return False # Wait for user to complete CAS login print(f"\n 请在浏览器中完成机构登录 ({idp_name})") print(" 登录成功后浏览器会自动跳转回文章页面，程序会自动检测\n") for i in range(100): - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return False try: title = page.title() url = page.url @@ -466,6 +584,7 @@ def _visible_browser_download( output_path: Path, config: dict[str, Any], publisher: str, + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Open a visible browser for full SSO login + direct PDF download. @@ -475,6 +594,8 @@ def _visible_browser_download( from .pdf_utils import is_pdf_file, success import base64 + if _cancelled(cancel_event): + return None idp_name = config.get("carsi_idp_name", "") idp_en = _IDP_MAP.get(idp_name, idp_name) @@ -484,7 +605,11 @@ def _visible_browser_download( search_selectors = sso_cfg["search_selectors"] pdf_paths = sso_cfg["pdf_paths"](doi) - with _visible_browser(config, publisher) as (context, page): + with _visible_browser( + config, + publisher, + cancel_event=cancel_event, + ) as (context, page): # For Elsevier DOIs, avoid linkinghub.elsevier.com by using direct URL if publisher == "Elsevier" and ("doi.org/" in article_url or "linkinghub" in article_url): @@ -501,11 +626,14 @@ def _visible_browser_download( # Navigate to article page log.info(f" [{publisher}] visible browser: opening {article_url[:60]}") + if _cancelled(cancel_event): + return None try: page.goto(article_url, wait_until="domcontentloaded", timeout=60000) - time.sleep(5) except Exception as exc: log.info(f" [{publisher}] page load warning: {exc}") + if _wait_or_cancel(cancel_event, 5): + return None # If stuck on linkinghub redirect, extract target and navigate directly url = page.url @@ -521,9 +649,12 @@ def _visible_browser_download( direct_urls.append(cell_url) direct_urls.append(f"https://www.sciencedirect.com/science/article/pii/{pii}") for direct_url in direct_urls: + if _cancelled(cancel_event): + return None try: page.goto(direct_url, wait_until="domcontentloaded", timeout=30000) - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None if "linkinghub" not in page.url: log.info(f" [{publisher}] navigated to {page.url[:60]}") break @@ -531,11 +662,15 @@ def _visible_browser_download( pass # Wait for Cloudflare challenge to resolve (visible browser can pass it) + from .network import is_cloudflare_challenge + for _cf_wait in range(12): - _cf_title = (page.title() or "").lower() - if any(_sig in _cf_title for _sig in ("just a moment", "attention required", "verify", "security check", "请稍候", "正在验证", "checking", "cloudflare")): + if _cancelled(cancel_event): + return None + if is_cloudflare_challenge(page.title() or ""): log.info(f" [{publisher}] Cloudflare challenge detected, waiting... ({_cf_wait+1}/12)") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None else: break else: @@ -553,23 +688,29 @@ def _visible_browser_download( pdf_fetched = False if not already_on_auth: log.info(f" [{publisher}] trying direct PDF fetch...") - fetch_result = _try_browser_fetch_pdf(page, pdf_paths) + fetch_result = _try_browser_fetch_pdf( + page, + pdf_paths, + cancel_event=cancel_event, + ) if fetch_result: pdf_bytes = fetch_result - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes via direct fetch") - if is_pdf_file(output_path): - return success(doi, output_path, f"{publisher}(Visible)") - pdf_fetched = True + if write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes via direct fetch") + if is_pdf_file(output_path): + return success(doi, output_path, f"{publisher}(Visible)") + pdf_fetched = True else: log.info(f" [{publisher}] direct PDF fetch failed, need SSO login") # SSO login needed if not already_on_auth and not pdf_fetched: log.info(f" [{publisher}] starting SSO login...") + if _cancelled(cancel_event): + return None page.evaluate(sso_link_js) - time.sleep(8) + if _wait_or_cancel(cancel_event, 8): + return None # SSO click may have navigated the page or opened a popup # Check all pages in the context for the SSO/IDP page @@ -590,10 +731,13 @@ def _visible_browser_download( # Search for institution on the SSO page for sel in search_selectors: + if _cancelled(cancel_event): + return None si = sso_page.query_selector(sel) if si: si.fill(idp_en) - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None clicked = sso_page.evaluate(f""" (name) => {{ const items = document.querySelectorAll('[class*="result"], [class*="suggestion"], [class*="federation"], li, a, button'); @@ -608,7 +752,8 @@ def _visible_browser_download( """, idp_en) if clicked: log.info(f" [{publisher}] selected institution '{idp_en}'") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None break # Wait for user to complete CAS login @@ -617,7 +762,8 @@ def _visible_browser_download( login_ok = False for i in range(100): - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None # Check ALL pages in context — SSO may happen in any tab any_auth = False for ctx_page in context.pages: @@ -640,7 +786,8 @@ def _visible_browser_download( return None log.info(f" [{publisher}] login successful, downloading PDF...") - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None # Save cookies for future use (all formats + bridge to browser-engine) try: @@ -651,24 +798,32 @@ def _visible_browser_download( # Navigate back to the article page and reload to pick up auth cookies log.info(f" [{publisher}] reloading article page with new auth cookies...") + if _cancelled(cancel_event): + return None try: page.goto(article_url, wait_until="domcontentloaded", timeout=30000) - time.sleep(5) except Exception: pass + if _wait_or_cancel(cancel_event, 5): + return None # Try downloading PDF via in-browser fetch (post-login) - fetch_result = _try_browser_fetch_pdf(page, pdf_paths) + fetch_result = _try_browser_fetch_pdf( + page, + pdf_paths, + cancel_event=cancel_event, + ) if fetch_result: pdf_bytes = fetch_result - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes via post-login fetch") - if is_pdf_file(output_path): - return success(doi, output_path, f"{publisher}(Visible)") + if write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes via post-login fetch") + if is_pdf_file(output_path): + return success(doi, output_path, f"{publisher}(Visible)") # Fallback: look for PDF link on the page and fetch it log.info(f" [{publisher}] fetch failed, trying page PDF links...") + if _cancelled(cancel_event): + return None pdf_link = page.evaluate(""" (() => { for (const a of document.querySelectorAll('a')) { @@ -687,11 +842,14 @@ def _visible_browser_download( if pdf_link and isinstance(pdf_link, str): log.info(f" [{publisher}] found PDF link: {pdf_link[:60]}") + if _cancelled(cancel_event): + return None try: page.goto(pdf_link, wait_until="domcontentloaded", timeout=30000) - time.sleep(3) except Exception: pass + if _wait_or_cancel(cancel_event, 3): + return None # Try fetch from current page context current_url = page.url @@ -699,21 +857,30 @@ def _visible_browser_download( parsed = urlparse(current_url) # Try absolute fetch abs_paths = [parsed.path] - fetch_result2 = _try_browser_fetch_pdf(page, abs_paths) + fetch_result2 = _try_browser_fetch_pdf( + page, + abs_paths, + cancel_event=cancel_event, + ) if fetch_result2: pdf_bytes = fetch_result2 - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes via PDF link") - if is_pdf_file(output_path): - return success(doi, output_path, f"{publisher}(Visible)") + if write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes via PDF link") + if is_pdf_file(output_path): + return success(doi, output_path, f"{publisher}(Visible)") log.info(f" [{publisher}] visible browser download failed") return None -def _try_browser_fetch_pdf(page: Any, paths: list[str]) -> bytes | None: +def _try_browser_fetch_pdf( + page: Any, + paths: list[str], + cancel_event: threading.Event | None = None, +) -> bytes | None: """Try fetching PDF via browser JS fetch for given paths. Returns PDF bytes or None.""" + if _cancelled(cancel_event): + return None import base64 paths_js = json.dumps(paths) pdf_b64 = page.evaluate(f""" @@ -737,6 +904,8 @@ def _try_browser_fetch_pdf(page: Any, paths: list[str]) -> bytes | None: }})() """) + if _cancelled(cancel_event): + return None if isinstance(pdf_b64, str) and pdf_b64.startswith("data:"): _, data = pdf_b64.split(",", 1) pdf_bytes = base64.b64decode(data) @@ -1244,6 +1413,7 @@ def _try_http_download( output_path: Path, config: dict[str, Any], cookies: dict[str, str] | None = None, + cancel_event: threading.Event | None = None, ) -> bool: """Try direct HTTP download of a PDF URL. Returns True on success.""" import requests @@ -1251,32 +1421,62 @@ def _try_http_download( from .pdf_utils import _response_looks_pdf, is_pdf_file from .sources.publishers import _write_pdf_atomic, _load_publisher_cookies + if _cancelled(cancel_event): + return False + session = None + resp = None try: - s = requests.Session() - s.trust_env = False - s.headers.update({"User-Agent": USER_AGENT}) - _load_publisher_cookies(s, config) + session = requests.Session() + session.trust_env = False + session.headers.update({"User-Agent": USER_AGENT}) + _load_publisher_cookies(session, config) if cookies: for name, value in cookies.items(): - s.cookies.set(name, value) + if _cancelled(cancel_event): + return False + session.cookies.set(name, value) - resp = s.get(pdf_url, timeout=20, stream=True, - headers={"Accept": "application/pdf,*/*"}, - allow_redirects=True) + if _cancelled(cancel_event): + return False + resp = session.get( + pdf_url, + timeout=20, + stream=True, + headers={"Accept": "application/pdf,*/*"}, + allow_redirects=True, + ) - if resp.status_code >= 400: + if _cancelled(cancel_event) or resp.status_code >= 400: return False iterator = resp.iter_content(chunk_size=8192) first = next(iterator, b"") + if _cancelled(cancel_event): + return False if not _response_looks_pdf(resp, first): return False - if not _write_pdf_atomic(output_path, first, iterator): + if not _write_pdf_atomic( + output_path, + first, + iterator, + cancel_event=cancel_event, + ): return False - return is_pdf_file(output_path) + return not _cancelled(cancel_event) and is_pdf_file(output_path) except Exception: return False + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass + if session is not None: + try: + session.close() + except Exception: + pass def _extract_pdf_from_page( @@ -1308,18 +1508,31 @@ def _extract_pdf_from_page( except Exception: continue - # Generic fallback: look for any link containing "pdf" (exclude supplements) + # Generic fallback: only accept same-site PDF links. Article reference lists + # commonly contain cross-publisher PDF URLs, which are not the article PDF. + from urllib.parse import urljoin, urlparse + + article_host = (urlparse(article_url).hostname or "").lower() for a in soup.find_all("a", href=True): href = a["href"].lower() + href_path = urlparse(href).path # Skip supplement/attachment links if any(skip in href for skip in ["/attachment/", "/cms/", "/mmc", "supplement", "supporting"]): continue - if "/pdf/" in href or href.endswith(".pdf") or "pdfdirect" in href: - if a["href"].startswith("http"): - return a["href"] - elif a["href"].startswith("/"): - from urllib.parse import urljoin - return urljoin(article_url, a["href"]) + if ( + "/pdf/" in href + or href.endswith(".pdf") + or href_path.endswith("/pdf") + or "pdfdirect" in href + ): + candidate = urljoin(article_url, a["href"]) + candidate_host = (urlparse(candidate).hostname or "").lower() + if candidate_host and ( + candidate_host == article_host + or candidate_host.endswith(f".{article_host}") + or article_host.endswith(f".{candidate_host}") + ): + return candidate return None @@ -1621,6 +1834,7 @@ def _browser_download( publisher: str, *, wait_for_loading: float = 0, + cancel_event: threading.Event | None = None, ) -> bool: """Navigate to article page via browser, find PDF link, download it.""" from .browser_engine import ( @@ -1631,12 +1845,23 @@ def _browser_download( from .pdf_utils import is_pdf_file _clear_error() + if _cancelled(cancel_event): + return False # Campus network fast-path: skip HTTP, go directly to CloakBrowser # Campus networks use IP authentication, so CloakBrowser can access directly if _is_campus_network(config) and is_available(config): log.info(f" [{publisher}] campus network detected, using CloakBrowser directly") - if download_pdf_via_browser(article_url, output_path, config): + if _cancelled(cancel_event): + return False + if download_pdf_via_browser( + article_url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return False from .pdf_utils import is_pdf_file, success if is_pdf_file(output_path): log.info(f" [{publisher}] campus network download succeeded") @@ -1646,7 +1871,16 @@ def _browser_download( # This avoids 15-30s browser startup overhead if _has_publisher_cookies(config) and _is_pdf_url(article_url): log.info(f" [{publisher}] trying HTTP with cookies first (fast-path)") - if _try_http_download(article_url, output_path, config): + if _cancelled(cancel_event): + return False + if _try_http_download( + article_url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return False from .pdf_utils import is_pdf_file, success if is_pdf_file(output_path): log.info(f" [{publisher}] HTTP download succeeded with cached cookies") @@ -1659,39 +1893,54 @@ def _browser_download( log.info(f" [{publisher}] browser download: {article_url[:80]}") - # Create tab to a lightweight page first, inject cookies, then navigate to target - tab_id = create_tab("https://www.google.com/", config, timeout=15.0) + # Start from a local blank document so cookie injection does not depend on an + # unrelated public site being reachable from the deployment network. + if _cancelled(cancel_event): + return False + tab_id = create_tab("about:blank", config, timeout=15.0) if not tab_id: return False try: # Inject saved publisher/CARSI cookies into the browser session + if _cancelled(cancel_event): + return False _inject_cookies_to_tab(tab_id, config, publisher) # Now navigate to the actual article page log.info(f" [{publisher}] navigating to {article_url[:80]}") + if _cancelled(cancel_event): + return False nav_ok = navigate_tab(tab_id, article_url, config, timeout=60.0) + if _cancelled(cancel_event): + return False if not nav_ok: log.info(f" [{publisher}] navigation failed, tab may be destroyed") _set_error("navigate_failed", "cloudflare_timeout") return False # Wait for loading pages (AIP/AVS) - if wait_for_loading > 0: - time.sleep(wait_for_loading) + if wait_for_loading > 0 and _wait_or_cancel(cancel_event, wait_for_loading): + return False # Wait for page to settle - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return False # Detect meta-refresh redirects (e.g., linkinghub.elsevier.com) # Wait longer if we're on a redirect page current_url = evaluate_js(tab_id, "window.location.href", config) or "" + if _cancelled(cancel_event): + return False if "linkinghub" in current_url or "retrieve/pii" in current_url: log.info(f" [{publisher}] on Elsevier redirect hub, waiting for meta-refresh...") # Wait for the browser to follow the meta-refresh chain for _ in range(5): - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return False new_url = evaluate_js(tab_id, "window.location.href", config) or "" + if _cancelled(cancel_event): + return False if new_url != current_url and "linkinghub" not in new_url and "retrieve/pii" not in new_url: log.info(f" [{publisher}] redirected to {new_url[:80]}") break @@ -1712,22 +1961,33 @@ def _browser_download( return null; })() """, config) + if _cancelled(cancel_event): + return False if redirect_url and isinstance(redirect_url, str) and redirect_url.startswith("http"): log.info(f" [{publisher}] manually following redirect to {redirect_url[:80]}") + if _cancelled(cancel_event): + return False navigate_tab(tab_id, redirect_url, config, timeout=30.0) - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return False # Wait for the final page to fully render - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return False # Get page HTML html = evaluate_js(tab_id, "document.documentElement.outerHTML", config) or "" + if _cancelled(cancel_event): + return False # Check for anti-bot challenges if _is_challenge_page(html): log.info(f" [{publisher}] challenge detected, waiting for auto-resolve...") - time.sleep(8) + if _wait_or_cancel(cancel_event, 8): + return False html = evaluate_js(tab_id, "document.documentElement.outerHTML", config) or "" + if _cancelled(cancel_event): + return False if _is_challenge_page(html): log.info(f" [{publisher}] challenge did not resolve") _set_error("cloudflare_blocked", "use_proxy_or_browser") @@ -1736,11 +1996,19 @@ def _browser_download( # Check for paywall AFTER challenge resolution if _detect_paywall(html): log.info(f" [{publisher}] paywall detected — trying institutional login...") - if _try_institutional_login(tab_id, config, publisher): + if _try_institutional_login( + tab_id, + config, + publisher, + cancel_event=cancel_event, + ): # Login succeeded, re-fetch page content - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return False html = evaluate_js(tab_id, "document.documentElement.outerHTML", config) or html current_url = evaluate_js(tab_id, "window.location.href", config) or article_url + if _cancelled(cancel_event): + return False if _detect_paywall(html): log.info(f" [{publisher}] still behind paywall after login") _set_error("paywall", "login_required") @@ -1751,11 +2019,15 @@ def _browser_download( # Also detect paywall by absence of PDF links (Cell Press pattern) if publisher == "Elsevier" and "cell.com" in str(evaluate_js(tab_id, "window.location.href", config) or ""): + if _cancelled(cancel_event): + return False has_showpdf = evaluate_js(tab_id, """ (() => { return document.querySelectorAll('a[href*="showPdf"], a[href*="pdfExtended"]').length; })() """, config) + if _cancelled(cancel_event): + return False # Use JS to check for institutional access links (HTML can be 500K+) has_institutional = evaluate_js(tab_id, """ (() => { @@ -1769,6 +2041,8 @@ def _browser_download( return false; })() """, config) + if _cancelled(cancel_event): + return False if has_institutional and not has_showpdf: log.info(f" [{publisher}] Cell Press paywall detected (institutional links present, no PDF links)") _set_error("paywall", "login_required") @@ -1777,19 +2051,29 @@ def _browser_download( # Elsevier-specific: detect crasolve if publisher == "Elsevier": current_url = evaluate_js(tab_id, "window.location.href", config) or article_url + if _cancelled(cancel_event): + return False if _is_elsevier_crasolve(str(current_url), html): log.info(f" [{publisher}] crasolve shell detected, waiting...") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return False html = evaluate_js(tab_id, "document.documentElement.outerHTML", config) or "" + if _cancelled(cancel_event): + return False # AIP-specific: check for loading page if publisher == "AIP" and _is_aip_loading_page(html): log.info(f" [{publisher}] loading page detected, waiting...") - time.sleep(10) + if _wait_or_cancel(cancel_event, 10): + return False html = evaluate_js(tab_id, "document.documentElement.outerHTML", config) or "" + if _cancelled(cancel_event): + return False # Get current URL after potential redirects current_url = evaluate_js(tab_id, "window.location.href", config) or article_url + if _cancelled(cancel_event): + return False # Try to extract PDF link from page pdf_url = _extract_pdf_from_page(html, str(current_url), publisher) @@ -1798,34 +2082,25 @@ def _browser_download( log.info(f" [{publisher}] found PDF link: {pdf_url[:80]}") # Strategy 1: Network response capture (navigate and intercept PDF at network layer) + if _cancelled(cancel_event): + return False captured = fetch_url(tab_id, pdf_url, config, timeout=30.0) + if _cancelled(cancel_event): + return False if captured and captured.get("data"): pdf_data = captured["data"] - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_data) + if not write_pdf_bytes_atomic(output_path, pdf_data, cancel_event): + return False log.info(f" [{publisher}] downloaded {len(pdf_data)} bytes via network capture") close_tab(tab_id, config) from .pdf_utils import is_pdf_file, success if is_pdf_file(output_path): return success(doi, output_path, f"{publisher}(Network)") - return True + output_path.unlink(missing_ok=True) + return False - # Strategy 2: Check any already-captured responses - captured_resps = get_captured_responses(tab_id, config, consume=True) - import base64 as _b64_net - for resp in captured_resps: - data = resp.get("dataBase64", "") - if data: - pdf_bytes = _b64_net.b64decode(data) - if pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes from captured response") - close_tab(tab_id, config) - from .pdf_utils import is_pdf_file, success - if is_pdf_file(output_path): - return success(doi, output_path, f"{publisher}(Network)") - return True + # Strategy 2: Clear response metadata retained by the tab listener + get_captured_responses(tab_id, config, consume=True) # Strategy 3: Try in-browser fetch for the PDF URL # This bypasses Cloudflare because the request comes from the browser context @@ -1841,6 +2116,8 @@ def _browser_download( fetch_paths.append(parsed.path.replace("/doi/epdf/", "/doi/pdfdirect/")) for fetch_path in fetch_paths: + if _cancelled(cancel_event): + return False log.info(f" [{publisher}] trying in-browser fetch {fetch_path[:60]}") pdf_b64 = evaluate_js(tab_id, f""" (async () => {{ @@ -1863,23 +2140,33 @@ def _browser_download( }} }})() """, config, timeout=45.0) + if _cancelled(cancel_event): + return False if isinstance(pdf_b64, str) and pdf_b64.startswith("data:"): header, data = pdf_b64.split(",", 1) pdf_bytes = _b64.b64decode(data) if pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) + if not write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + return False log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes via in-browser fetch") close_tab(tab_id, config) from .pdf_utils import is_pdf_file, success if is_pdf_file(output_path): return success(doi, output_path, f"{publisher}(Browser)") - return True + output_path.unlink(missing_ok=True) + return False elif isinstance(pdf_b64, str) and "status:403" in pdf_b64: log.info(f" [{publisher}] PDF fetch returned 403 — trying institutional login...") - if _try_institutional_login(tab_id, config, publisher): + if _try_institutional_login( + tab_id, + config, + publisher, + cancel_event=cancel_event, + ): # Retry fetch after login + if _cancelled(cancel_event): + return False pdf_b64_retry = evaluate_js(tab_id, f""" (async () => {{ try {{ @@ -1901,36 +2188,63 @@ def _browser_download( }} }})() """, config, timeout=45.0) + if _cancelled(cancel_event): + return False if isinstance(pdf_b64_retry, str) and pdf_b64_retry.startswith("data:"): _, data = pdf_b64_retry.split(",", 1) pdf_bytes = _b64.b64decode(data) if pdf_bytes[:5] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) + if not write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + return False log.info(f" [{publisher}] downloaded {len(pdf_bytes)} bytes after institutional login") close_tab(tab_id, config) from .pdf_utils import is_pdf_file, success if is_pdf_file(output_path): return success(doi, output_path, f"{publisher}(Institutional)") - return True + output_path.unlink(missing_ok=True) + return False _set_error("paywall", "login_required") elif isinstance(pdf_b64, str) and "status:401" in pdf_b64: log.info(f" [{publisher}] PDF fetch returned 401 — trying institutional login...") - if not _try_institutional_login(tab_id, config, publisher): + if not _try_institutional_login( + tab_id, + config, + publisher, + cancel_event=cancel_event, + ): _set_error("paywall", "login_required") elif isinstance(pdf_b64, str) and pdf_b64.startswith("ct:text/html"): log.info(f" [{publisher}] PDF fetch returned HTML — trying institutional login...") - if not _try_institutional_login(tab_id, config, publisher): + if not _try_institutional_login( + tab_id, + config, + publisher, + cancel_event=cancel_event, + ): _set_error("paywall", "login_required") close_tab(tab_id, config) # Fall back to HTTP download - if _try_http_download(pdf_url, output_path, config): + if _cancelled(cancel_event): + return False + if _try_http_download( + pdf_url, + output_path, + config, + cancel_event=cancel_event, + ): return True # Fall back to full browser download - return download_pdf_via_browser(pdf_url, output_path, config) + if _cancelled(cancel_event): + return False + return download_pdf_via_browser( + pdf_url, + output_path, + config, + cancel_event=cancel_event, + ) # For Elsevier/Cell Press: try specific PDF link patterns if publisher == "Elsevier": @@ -1950,13 +2264,29 @@ def _browser_download( return null; })() """, config) + if _cancelled(cancel_event): + return False if pdf_url and isinstance(pdf_url, str): log.info(f" [{publisher}] Elseviewer PDF button: {pdf_url[:80]}") close_tab(tab_id, config) - if _try_http_download(pdf_url, output_path, config): + if _cancelled(cancel_event): + return False + if _try_http_download( + pdf_url, + output_path, + config, + cancel_event=cancel_event, + ): return True - return download_pdf_via_browser(pdf_url, output_path, config) + if _cancelled(cancel_event): + return False + return download_pdf_via_browser( + pdf_url, + output_path, + config, + cancel_event=cancel_event, + ) # For Wiley: specifically look for pdfdirect if publisher == "Wiley": @@ -1975,13 +2305,29 @@ def _browser_download( return null; })() """, config) + if _cancelled(cancel_event): + return False if pdf_url and isinstance(pdf_url, str): log.info(f" [{publisher}] Wiley PDFDirect: {pdf_url[:80]}") close_tab(tab_id, config) - if _try_http_download(pdf_url, output_path, config): + if _cancelled(cancel_event): + return False + if _try_http_download( + pdf_url, + output_path, + config, + cancel_event=cancel_event, + ): return True - return download_pdf_via_browser(pdf_url, output_path, config) + if _cancelled(cancel_event): + return False + return download_pdf_via_browser( + pdf_url, + output_path, + config, + cancel_event=cancel_event, + ) # For IEEE: look for stamp URL if publisher == "IEEE": @@ -1995,11 +2341,20 @@ def _browser_download( return null; })() """, config) + if _cancelled(cancel_event): + return False if pdf_url and isinstance(pdf_url, str): log.info(f" [{publisher}] IEEE stamp: {pdf_url[:80]}") close_tab(tab_id, config) - return download_pdf_via_browser(pdf_url, output_path, config) + if _cancelled(cancel_event): + return False + return download_pdf_via_browser( + pdf_url, + output_path, + config, + cancel_event=cancel_event, + ) log.info(f" [{publisher}] no PDF link found on page") if not _last_error_type: @@ -2015,11 +2370,28 @@ def _browser_download_with_fallback( output_path: Path, config: dict[str, Any], publisher: str, + *, + wait_for_loading: float = 0, + cancel_event: threading.Event | None = None, ) -> bool: """Try headless browser download, fallback to visible browser if paywall detected.""" - result = _browser_download(doi, article_url, output_path, config, publisher) + if _cancelled(cancel_event): + return False + result = _browser_download( + doi, + article_url, + output_path, + config, + publisher, + wait_for_loading=wait_for_loading, + cancel_event=cancel_event, + ) + if _cancelled(cancel_event): + return False if result: return True + if _cancelled(cancel_event): + return False err_type, err_action = get_last_error() @@ -2047,11 +2419,21 @@ def _browser_download_with_fallback( elif err_type in ("paywall", "navigate_failed", "cloudflare_blocked") and config.get("carsi_idp_name"): log.info(f" [{publisher}] headless failed ({err_type}), trying visible browser fallback...") try: - vbd_result = _visible_browser_download(doi, article_url, output_path, config, publisher) + vbd_result = _visible_browser_download( + doi, + article_url, + output_path, + config, + publisher, + cancel_event=cancel_event, + ) + if _cancelled(cancel_event): + return False if vbd_result: return True except Exception as e: - log.info(f" [{publisher}] visible browser fallback error: {e}") + if not _cancelled(cancel_event): + log.info(f" [{publisher}] visible browser fallback error: {e}") return False @@ -2283,22 +2665,22 @@ def _save_elsevier_pdf_content( source: str, *, reject_single_page: bool = False, + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: from .pdf_utils import is_pdf_file, success - if len(content) < config.get("min_pdf_size_bytes", 10000): + if _cancelled(cancel_event) or len(content) < config.get("min_pdf_size_bytes", 10000): log.info(f" [ElsevierAPI] response too small ({len(content)} bytes)") return None if reject_single_page: page_count = _elsevier_pdf_page_count(content) if page_count == 1: - output_path.unlink(missing_ok=True) log.info(f" [ElsevierAPI] direct PDF is a 1-page preview for {doi}") return None - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(content) - if is_pdf_file(output_path): + if not write_pdf_bytes_atomic(output_path, content, cancel_event): + return None + if not _cancelled(cancel_event) and is_pdf_file(output_path): log.info(f" [ElsevierAPI] downloaded {len(content)} bytes for {doi}") return success(doi, output_path, source) @@ -2316,10 +2698,14 @@ def _try_elsevier_object_pdf_from_xml( first_resp: Any, *, full_xml_already_tried: bool = False, + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: import urllib.parse + if _cancelled(cancel_event): + return None xml_resp = None + owns_xml_resp = False if ( first_resp is not None and getattr(first_resp, "status_code", 0) == 200 @@ -2332,6 +2718,9 @@ def _try_elsevier_object_pdf_from_xml( xml_headers["Accept"] = "application/xml" param_options = [None] if full_xml_already_tried else [{"view": "FULL"}, None] for params in param_options: + if _cancelled(cancel_event): + return None + candidate = None try: request_kwargs: dict[str, Any] = { "headers": xml_headers, @@ -2341,89 +2730,120 @@ def _try_elsevier_object_pdf_from_xml( if params: request_kwargs["params"] = params candidate = session.get(article_url, **request_kwargs) + if _cancelled(cancel_event): + with contextlib.suppress(Exception): + candidate.close() + candidate = None + return None except Exception as exc: log.info(f" [ElsevierAPI] XML request failed: {exc}") continue - if getattr(candidate, "status_code", 0) != 200: - view_name = "FULL XML" if params else "XML" - log.info( - f" [ElsevierAPI] {view_name} HTTP " - f"{getattr(candidate, 'status_code', 0)} for {doi}" - ) - continue - if not _elsevier_response_is_xml(candidate): - content_type = _elsevier_header( - getattr(candidate, "headers", {}), - "content-type", - ) - view_name = "FULL XML" if params else "XML" - log.info( - f" [ElsevierAPI] {view_name} request returned non-XML " - f"({content_type[:50]})" - ) - continue + try: + if getattr(candidate, "status_code", 0) != 200: + view_name = "FULL XML" if params else "XML" + log.info( + f" [ElsevierAPI] {view_name} HTTP " + f"{getattr(candidate, 'status_code', 0)} for {doi}" + ) + continue + if not _elsevier_response_is_xml(candidate): + content_type = _elsevier_header( + getattr(candidate, "headers", {}), + "content-type", + ) + view_name = "FULL XML" if params else "XML" + log.info( + f" [ElsevierAPI] {view_name} request returned non-XML " + f"({content_type[:50]})" + ) + continue - xml_resp = candidate - break + xml_resp = candidate + owns_xml_resp = True + candidate = None + break + finally: + if candidate is not None: + with contextlib.suppress(Exception): + candidate.close() if xml_resp is None: return None - eids = _extract_elsevier_pdf_attachment_eids(_elsevier_response_text(xml_resp)) - if not eids: - log.info(f" [ElsevierAPI] XML has no PDF attachment EID for {doi}") - return None + try: + eids = _extract_elsevier_pdf_attachment_eids(_elsevier_response_text(xml_resp)) + if not eids: + log.info(f" [ElsevierAPI] XML has no PDF attachment EID for {doi}") + return None - object_headers = dict(headers) - object_headers["Accept"] = "application/pdf" - for eid in eids: - object_url = ( - "https://api.elsevier.com/content/object/eid/" - f"{urllib.parse.quote(eid, safe='')}" - ) - try: - object_resp = session.get( - object_url, - headers=object_headers, - timeout=30, - allow_redirects=True, + object_headers = dict(headers) + object_headers["Accept"] = "application/pdf" + for eid in eids: + if _cancelled(cancel_event): + return None + object_url = ( + "https://api.elsevier.com/content/object/eid/" + f"{urllib.parse.quote(eid, safe='')}" ) - except Exception as exc: - log.info(f" [ElsevierAPI] object {eid} request failed: {exc}") - continue + object_resp = None + try: + try: + object_resp = session.get( + object_url, + headers=object_headers, + timeout=30, + allow_redirects=True, + ) + if _cancelled(cancel_event): + continue + except Exception as exc: + log.info(f" [ElsevierAPI] object {eid} request failed: {exc}") + continue - if getattr(object_resp, "status_code", 0) != 200: - log.info( - f" [ElsevierAPI] object {eid} HTTP " - f"{getattr(object_resp, 'status_code', 0)}" - ) - continue - if not _elsevier_response_is_pdf(object_resp): - content_type = _elsevier_header( - getattr(object_resp, "headers", {}), - "content-type", - ) - log.info(f" [ElsevierAPI] object {eid} returned non-PDF ({content_type[:50]})") - continue + if getattr(object_resp, "status_code", 0) != 200: + log.info( + f" [ElsevierAPI] object {eid} HTTP " + f"{getattr(object_resp, 'status_code', 0)}" + ) + continue + if not _elsevier_response_is_pdf(object_resp): + content_type = _elsevier_header( + getattr(object_resp, "headers", {}), + "content-type", + ) + log.info(f" [ElsevierAPI] object {eid} returned non-PDF ({content_type[:50]})") + continue - result = _save_elsevier_pdf_content( - doi, - output_path, - getattr(object_resp, "content", b"") or b"", - config, - "ElsevierAPI", - reject_single_page=True, - ) - if result: - log.info(f" [ElsevierAPI] downloaded object PDF via attachment EID {eid}") - return result + result = _save_elsevier_pdf_content( + doi, + output_path, + getattr(object_resp, "content", b"") or b"", + config, + "ElsevierAPI", + reject_single_page=True, + cancel_event=cancel_event, + ) + if result: + log.info(f" [ElsevierAPI] downloaded object PDF via attachment EID {eid}") + return result + finally: + if object_resp is not None: + with contextlib.suppress(Exception): + object_resp.close() - return None + return None + finally: + if owns_xml_resp and xml_resp is not None: + with contextlib.suppress(Exception): + xml_resp.close() def try_elsevier_api( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Download Elsevier/ScienceDirect PDF via Article Retrieval API. @@ -2439,6 +2859,8 @@ def try_elsevier_api( elsevier_api_key — personal or institutional API key (required) elsevier_insttoken — institutional token for campus-level access """ + if _cancelled(cancel_event): + return None api_key = config.get("elsevier_api_key", "") if not api_key: log.info(f" [ElsevierAPI] no API key configured, skipping. " @@ -2466,43 +2888,69 @@ def try_elsevier_api( import requests for route_name, proxies in route_options: + if _cancelled(cancel_event): + return None + session = None + resp = None try: - session = requests.Session() - session.trust_env = False - if proxies: - session.proxies = proxies - - log.info(f" [ElsevierAPI] trying {route_name} route") - xml_headers = dict(headers) - xml_headers["Accept"] = "application/xml" - resp = session.get( - url, - headers=xml_headers, - params={"view": "FULL"}, - timeout=30, - allow_redirects=True, - ) - except Exception as e: - log.info(f" [ElsevierAPI] {route_name} FULL XML request failed: {e}") - continue + try: + session = requests.Session() + session.trust_env = False + if proxies: + session.proxies = proxies + + log.info(f" [ElsevierAPI] trying {route_name} route") + xml_headers = dict(headers) + xml_headers["Accept"] = "application/xml" + resp = session.get( + url, + headers=xml_headers, + params={"view": "FULL"}, + timeout=30, + allow_redirects=True, + ) + if _cancelled(cancel_event): + continue + except Exception as e: + log.info(f" [ElsevierAPI] {route_name} FULL XML request failed: {e}") + continue - if resp is not None and resp.status_code != 200: - if resp.status_code in (403, 429): - log.info( - f" [ElsevierAPI] {route_name} HTTP {resp.status_code}; " - "API access may be unavailable or rate limited" + if resp is not None and resp.status_code != 200: + if resp.status_code in (403, 429): + log.info( + f" [ElsevierAPI] {route_name} HTTP {resp.status_code}; " + "API access may be unavailable or rate limited" + ) + else: + log.info(f" [ElsevierAPI] {route_name} HTTP {resp.status_code} for {doi}") + + if resp is not None and resp.status_code == 200 and _elsevier_response_is_pdf(resp): + result = _save_elsevier_pdf_content( + doi, + output_path, + resp.content, + config, + "ElsevierAPI", + reject_single_page=True, + cancel_event=cancel_event, ) - else: - log.info(f" [ElsevierAPI] {route_name} HTTP {resp.status_code} for {doi}") + if result: + try: + _persist_api_cookies(session, config) + except Exception: + pass + return result - if resp is not None and resp.status_code == 200 and _elsevier_response_is_pdf(resp): - result = _save_elsevier_pdf_content( + result = _try_elsevier_object_pdf_from_xml( doi, output_path, - resp.content, config, - "ElsevierAPI", - reject_single_page=True, + session, + url, + headers, + resp, + full_xml_already_tried=True, + cancel_event=cancel_event, ) if result: try: @@ -2511,35 +2959,25 @@ def try_elsevier_api( pass return result - result = _try_elsevier_object_pdf_from_xml( - doi, - output_path, - config, - session, - url, - headers, - resp, - full_xml_already_tried=True, - ) - if result: + content_type = _elsevier_header( + getattr(resp, "headers", {}) if resp is not None else {}, + "content-type", + ) + log.info( + f" [ElsevierAPI] {route_name} non-PDF response " + f"({content_type[:50]}), trying next route if available" + ) try: _persist_api_cookies(session, config) - except Exception: - pass - return result - - content_type = _elsevier_header( - getattr(resp, "headers", {}) if resp is not None else {}, - "content-type", - ) - log.info( - f" [ElsevierAPI] {route_name} non-PDF response " - f"({content_type[:50]}), trying next route if available" - ) - try: - _persist_api_cookies(session, config) - except Exception as e: - log.info(f" [ElsevierAPI] cookie persist failed: {e}") + except Exception as e: + log.info(f" [ElsevierAPI] cookie persist failed: {e}") + finally: + if resp is not None: + with contextlib.suppress(Exception): + resp.close() + if session is not None: + with contextlib.suppress(Exception): + session.close() return None # ============================================================ @@ -2547,19 +2985,33 @@ def try_elsevier_api( # ============================================================ def try_elsevier_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Elsevier/ScienceDirect/Cell Press browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success from .browser_engine import is_available as browser_available, download_pdf_via_browser + if _cancelled(cancel_event): + return None # Campus network fast-path: skip HTTP, go directly to CloakBrowser if _is_campus_network(config) and browser_available(config): cell_url = _build_cell_press_url(doi) if cell_url: log.info(f" [Elsevier] campus network detected, trying CloakBrowser directly: {cell_url[:80]}") - if download_pdf_via_browser(cell_url, output_path, config): + if _cancelled(cancel_event): + return None + if download_pdf_via_browser( + cell_url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): log.info(f" [Elsevier] campus network download succeeded") return success(doi, output_path, "CellPress(Campus)") @@ -2573,7 +3025,16 @@ def try_elsevier_browser( pii = cell_url.split("/abstract/")[-1] show_pdf_url = f"https://www.cell.com/action/showPdf?pii={pii}" log.info(f" [Elsevier] trying Cell Press showPdf HTTP with cookies: {show_pdf_url[:80]}") - if _try_http_download(show_pdf_url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + show_pdf_url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): log.info(f" [Elsevier] Cell Press HTTP download succeeded with cached cookies") return success(doi, output_path, "CellPress(HTTP)") @@ -2584,13 +3045,25 @@ def try_elsevier_browser( pii = cell_url.split("/abstract/")[-1] show_pdf_url = f"https://www.cell.com/action/showPdf?pii={pii}" log.info(f" [CellPress] trying showPdf via network capture: {show_pdf_url[:80]}") - result = _cell_press_showpdf_download(show_pdf_url, output_path, config) + result = _cell_press_showpdf_download( + show_pdf_url, + output_path, + config, + cancel_event=cancel_event, + ) if result and is_pdf_file(output_path): return result # If showPdf failed, navigate to the abstract page directly log.info(f" [CellPress] trying abstract page: {cell_url[:80]}") - if _browser_download_with_fallback(doi, cell_url, output_path, config, "Elsevier"): + if _browser_download_with_fallback( + doi, + cell_url, + output_path, + config, + "Elsevier", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "CellPress(Browser)") @@ -2598,8 +3071,19 @@ def try_elsevier_browser( # Only try if we haven't already detected a paywall if not _last_error_type: # Resolve DOI to direct sciencedirect.com URL (avoid linkinghub) + if _cancelled(cancel_event): + return None article_url = _resolve_elsevier_pii(doi, config) or f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "Elsevier"): + if _cancelled(cancel_event): + return None + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "Elsevier", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "Elsevier(Browser)") return None @@ -2609,26 +3093,33 @@ def _cell_press_showpdf_download( show_pdf_url: str, output_path: Path, config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Download a Cell Press showPdf URL via browser network capture.""" from .browser_engine import create_tab, close_tab, fetch_url from .pdf_utils import success - tab_id = create_tab("https://example.com", config, timeout=15.0) + if _cancelled(cancel_event): + return None + tab_id = create_tab("about:blank", config, timeout=15.0) if not tab_id: return None try: + if _cancelled(cancel_event): + return None _inject_cookies_to_tab(tab_id, config, "Elsevier") + if _cancelled(cancel_event): + return None result = fetch_url(tab_id, show_pdf_url, config, timeout=30.0) - if result and result.get("data"): + if not _cancelled(cancel_event) and result and result.get("data"): pdf_bytes = result["data"] - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - log.info(f" [CellPress] downloaded {len(pdf_bytes)} bytes via showPdf") - return success(show_pdf_url, output_path, "CellPress(Network)") + if write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + log.info(f" [CellPress] downloaded {len(pdf_bytes)} bytes via showPdf") + return success(show_pdf_url, output_path, "CellPress(Network)") except Exception as e: - log.info(f" [CellPress] showPdf failed: {e}") + if not _cancelled(cancel_event): + log.info(f" [CellPress] showPdf failed: {e}") finally: close_tab(tab_id, config) return None @@ -2684,6 +3175,7 @@ def _build_cell_press_url(doi: str) -> str | None: def _get_elsevier_pii(doi: str) -> str | None: """Get PII (Publisher Item Identifier) from Crossref API.""" + resp = None try: import requests from .network import USER_AGENT @@ -2701,6 +3193,10 @@ def _get_elsevier_pii(doi: str) -> str | None: return aid except Exception: pass + finally: + if resp is not None: + with contextlib.suppress(Exception): + resp.close() return None @@ -2711,6 +3207,8 @@ def _resolve_elsevier_pii(doi: str, config: dict[str, Any]) -> str | None: Returns a direct cell.com/sciencedirect.com URL, or None. """ import re + s = None + resp = None try: import requests from .network import USER_AGENT @@ -2741,244 +3239,498 @@ def _resolve_elsevier_pii(doi: str, config: dict[str, Any]) -> str | None: return None except Exception: return None + finally: + if resp is not None: + with contextlib.suppress(Exception): + resp.close() + if s is not None: + with contextlib.suppress(Exception): + s.close() def try_wiley_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Wiley browser strategy with PDFDirect.""" from .pdf_utils import is_pdf_file from .pdf_utils import success from .browser_engine import is_available as browser_available, download_pdf_via_browser + if _cancelled(cancel_event): + return None # Campus network fast-path: skip HTTP, go directly to CloakBrowser if _is_campus_network(config) and browser_available(config): article_url = f"https://doi.org/{doi}" log.info(f" [Wiley] campus network detected, trying CloakBrowser directly") - if download_pdf_via_browser(article_url, output_path, config): + if _cancelled(cancel_event): + return None + if download_pdf_via_browser( + article_url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): log.info(f" [Wiley] campus network download succeeded") return success(doi, output_path, "Wiley(Campus)") # Try direct PDF URLs first for url in _direct_pdf_urls(doi, "Wiley"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "Wiley(PDFDirect)") # Fall back to browser — use direct Wiley URL to avoid slow DOI redirect article_url = f"https://onlinelibrary.wiley.com/doi/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "Wiley"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "Wiley", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "Wiley(Browser)") return None def try_ieee_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """IEEE browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "IEEE"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "IEEE", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "IEEE(Browser)") return None def try_acs_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """ACS browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None # Try direct PDF URL first for url in _direct_pdf_urls(doi, "ACS"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "ACS(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "ACS"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "ACS", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "ACS(Browser)") return None def try_rsc_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """RSC browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None for url in _direct_pdf_urls(doi, "RSC"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "RSC(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "RSC"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "RSC", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "RSC(Browser)") return None def try_aip_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """AIP browser strategy with loading page wait.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "AIP", wait_for_loading=10): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "AIP", + wait_for_loading=10, + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "AIP(Browser)") return None def try_springer_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Springer browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None # Try direct PDF URL first for url in _direct_pdf_urls(doi, "Springer"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "Springer(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "Springer"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "Springer", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "Springer(Browser)") return None def try_aps_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """APS (Physical Review) browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None for url in _direct_pdf_urls(doi, "APS"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "APS(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "APS"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "APS", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "APS(Browser)") return None def try_tandfonline_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Taylor & Francis browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None for url in _direct_pdf_urls(doi, "Tandfonline"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "T&F(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "Tandfonline"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "Tandfonline", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "T&F(Browser)") return None def try_iop_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """IOP browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None for url in _direct_pdf_urls(doi, "IOP"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "IOP(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "IOP"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "IOP", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "IOP(Browser)") return None def try_oxford_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Oxford Academic browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None for url in _direct_pdf_urls(doi, "Oxford"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "Oxford(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "Oxford"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "Oxford", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "Oxford(Browser)") return None def try_acm_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """ACM browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None for url in _direct_pdf_urls(doi, "ACM"): - if _try_http_download(url, output_path, config): + if _cancelled(cancel_event): + return None + if _try_http_download( + url, + output_path, + config, + cancel_event=cancel_event, + ): + if _cancelled(cancel_event): + return None if is_pdf_file(output_path): return success(doi, output_path, "ACM(Direct)") article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "ACM"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "ACM", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "ACM(Browser)") return None def try_nature_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Nature browser strategy (fallback when direct fails).""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "Nature"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "Nature", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "Nature(Browser)") return None def try_science_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Science/AAAS browser strategy.""" from .pdf_utils import is_pdf_file from .pdf_utils import success + if _cancelled(cancel_event): + return None article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, "Science"): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + "Science", + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, "Science(Browser)") return None @@ -3122,17 +3874,29 @@ def try_copernicus_direct( # ============================================================ def try_generic_browser( - doi: str, output_path: Path, config: dict[str, Any], + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Generic browser fallback for unknown publishers.""" from .pdf_utils import is_pdf_file from .sources.publishers import get_publisher from .pdf_utils import success + if _cancelled(cancel_event): + return None publisher = get_publisher(doi) or "Unknown" article_url = f"https://doi.org/{doi}" - if _browser_download_with_fallback(doi, article_url, output_path, config, publisher): + if _browser_download_with_fallback( + doi, + article_url, + output_path, + config, + publisher, + cancel_event=cancel_event, + ): if is_pdf_file(output_path): return success(doi, output_path, f"{publisher}(Browser)") return None diff --git a/src/scansci_pdf/server.py b/src/scansci_pdf/server.py index 2b39a6a..f691165 100644 --- a/src/scansci_pdf/server.py +++ b/src/scansci_pdf/server.py @@ -411,23 +411,65 @@ def scansci_pdf_elsevier_setup(test: bool = False) -> str: import requests from .network import USER_AGENT try: - s = requests.Session() - s.trust_env = False proxy = config.get("network_proxy", "") + route_options: list[ + tuple[str, dict[str, str] | None] + ] = [("direct", None)] if proxy: - s.proxies = {"http": proxy, "https": proxy} - resp = s.get( - "https://api.elsevier.com/content/serial/title", - headers={"Accept": "application/json", "X-ELS-APIKey": api_key, "User-Agent": USER_AGENT}, - params={"count": 1}, - timeout=15, - ) - if resp.status_code == 200: + route_options.append(("configured_proxy", {"http": proxy, "https": proxy})) + + response_status = None + route_used = "" + last_error = "" + for route_name, proxies in route_options: + session = None + candidate = None + try: + session = requests.Session() + session.trust_env = False + if proxies: + session.proxies = proxies + candidate = session.get( + "https://api.elsevier.com/content/serial/title", + headers={ + "Accept": "application/json", + "X-ELS-APIKey": api_key, + "User-Agent": USER_AGENT, + }, + params={"count": 1}, + timeout=15, + ) + response_status = candidate.status_code + route_used = route_name + if candidate.status_code == 200: + break + except Exception as route_exc: + last_error = str(route_exc) + finally: + if candidate is not None: + try: + candidate.close() + except Exception: + pass + if session is not None: + try: + session.close() + except Exception: + pass + + if response_status is None: + raise RuntimeError(last_error or "no Elsevier API response") + + result["route"] = route_used + if response_status == 200: result["test"] = "passed" - result["message"] += " API Key 验证有效！ScienceDirect 论文可直接 API 下载。" + result["message"] += ( + " API Key 验证有效。闭源全文还需要机构订阅/IP entitlement；" + "下载时会优先走 direct route，再回退配置代理。" + ) else: result["test"] = "failed" - result["message"] += f" API Key 验证失败（HTTP {resp.status_code}），请检查 key 是否正确。" + result["message"] += f" API Key 验证失败（HTTP {response_status}），请检查 key 是否正确。" except Exception as e: result["test"] = "error" result["message"] += f" 验证请求失败: {e}" @@ -746,14 +788,17 @@ def scansci_pdf_carsi_login(publisher: str | None = None) -> str: target_publisher = publisher or "sciencedirect" client = CARSIClient(config) - if target_publisher not in client._publisher_configs: - available = list(client._publisher_configs.keys()) - return json.dumps({"success": False, "error": f"Unknown publisher: {target_publisher}", "available": available}) + try: + if target_publisher not in client._publisher_configs: + available = list(client._publisher_configs.keys()) + return json.dumps({"success": False, "error": f"Unknown publisher: {target_publisher}", "available": available}) - ok = client.login(target_publisher) - if ok: - return json.dumps({"success": True, "message": f"CARSI login successful for {target_publisher}.", "idp": idp_name}) - return json.dumps({"success": False, "error": "Login failed or timed out. Make sure Chrome is installed."}) + ok = client.login(target_publisher) + if ok: + return json.dumps({"success": True, "message": f"CARSI login successful for {target_publisher}.", "idp": idp_name}) + return json.dumps({"success": False, "error": "Login failed or timed out. Make sure Chrome is installed."}) + finally: + client.close() @mcp_app.tool() @@ -769,21 +814,24 @@ def scansci_pdf_carsi_status() -> str: return json.dumps({"carsi_enabled": False, "message": "CARSI not enabled."}) client = CARSIClient(config) - publishers = {} - for pub_key in client._publisher_configs: - cookie_file = client._cookie_path(pub_key) - has_cookies = cookie_file.exists() - publishers[pub_key] = { - "has_cookies": has_cookies, - "cookie_file": str(cookie_file), - } + try: + publishers = {} + for pub_key in client._publisher_configs: + cookie_file = client._cookie_path(pub_key) + has_cookies = cookie_file.exists() + publishers[pub_key] = { + "has_cookies": has_cookies, + "cookie_file": str(cookie_file), + } - return json.dumps({ - "carsi_enabled": True, - "carsi_idp_name": idp_name, - "hint": f"当前学校: {idp_name}。如需更换，运行 scansci_pdf_config_set key=carsi_idp_name value=新学校名称" if idp_name else "未设置学校。运行 scansci_pdf_config_set key=carsi_idp_name value=你的学校名称", - "publishers": publishers, - }, ensure_ascii=False) + return json.dumps({ + "carsi_enabled": True, + "carsi_idp_name": idp_name, + "hint": f"当前学校: {idp_name}。如需更换，运行 scansci_pdf_config_set key=carsi_idp_name value=新学校名称" if idp_name else "未设置学校。运行 scansci_pdf_config_set key=carsi_idp_name value=你的学校名称", + "publishers": publishers, + }, ensure_ascii=False) + finally: + client.close() @mcp_app.tool() diff --git a/src/scansci_pdf/sources/carsi.py b/src/scansci_pdf/sources/carsi.py index 45cc6e6..bbb8e79 100644 --- a/src/scansci_pdf/sources/carsi.py +++ b/src/scansci_pdf/sources/carsi.py @@ -6,6 +6,7 @@ from __future__ import annotations +import contextlib import json import os import re @@ -20,6 +21,7 @@ from ..config import DATA_DIR from ..log import get_logger +from ..pdf_utils import write_pdf_bytes_atomic from ..publisher_strategies import ( _IDP_MAP, _AUTH_KEYWORDS, @@ -31,6 +33,37 @@ log = get_logger() + +def _cancelled(cancel_event: threading.Event | None) -> bool: + return cancel_event is not None and cancel_event.is_set() + + +def _wait_or_cancel( + cancel_event: threading.Event | None, + timeout: float, +) -> bool: + if cancel_event is None: + time.sleep(max(0.0, timeout)) + return False + return cancel_event.wait(max(0.0, timeout)) + + +@contextlib.contextmanager +def _cancelable_lock( + lock: threading.Lock, + cancel_event: threading.Event | None, +): + acquired = False + while not acquired: + if _cancelled(cancel_event): + yield False + return + acquired = lock.acquire(timeout=0.1) + try: + yield True + finally: + lock.release() + _PUBLISHER_CONFIGS_FILE = DATA_DIR / "publisher_carsi.json" _PKG_DATA_DIR = Path(__file__).resolve().parent.parent / "data" _PKG_PUBLISHER_CONFIGS_FILE = _PKG_DATA_DIR / "publisher_carsi.json" @@ -126,12 +159,31 @@ def fetch(self, url: str, **kwargs) -> requests.Response | None: log.warning(f" [CARSI] Fetch failed: {e}") return None - def download_via_browser(self, doi: str, article_url: str, output_path: Path) -> dict[str, Any] | None: + def download_via_browser( + self, + doi: str, + article_url: str, + output_path: Path, + cancel_event: threading.Event | None = None, + ) -> dict[str, Any] | None: """Download PDF via CloakBrowser with CARSI auth.""" - return self._download_via_cloakbrowser(doi, article_url, output_path) - - def _download_via_cloakbrowser(self, doi: str, article_url: str, output_path: Path) -> dict[str, Any] | None: + return self._download_via_cloakbrowser( + doi, + article_url, + output_path, + cancel_event=cancel_event, + ) + + def _download_via_cloakbrowser( + self, + doi: str, + article_url: str, + output_path: Path, + cancel_event: threading.Event | None = None, + ) -> dict[str, Any] | None: """Download PDF via CloakBrowser with CARSI auth. Single session: login + download.""" + if _cancelled(cancel_event): + return None publisher = detect_publisher(article_url) if not publisher: return None @@ -140,8 +192,11 @@ def _download_via_cloakbrowser(self, doi: str, article_url: str, output_path: Pa return None try: + from ..cloakbrowser_compat import prepare_cloakbrowser_runtime + + prepare_cloakbrowser_runtime() from cloakbrowser import launch # noqa: F401 - except ImportError: + except Exception: log.info(" [CARSI-Browser] cloakbrowser not installed") return None @@ -155,18 +210,34 @@ def _download_via_cloakbrowser(self, doi: str, article_url: str, output_path: Pa from ..pdf_utils import is_pdf_file, success as _success # Serialize browser opens across threads — only one browser at a time - with self._login_lock: + with _cancelable_lock(self._login_lock, cancel_event) as lock_acquired: + if not lock_acquired: + return None log.info(f" [CARSI-Browser] Opening browser for {publisher}...") + page = None + captured_pdf: list[bytes] = [] + capture_lock = threading.Lock() + capture_in_progress = False + on_response = None try: from ..publisher_strategies import _visible_browser, _save_all_cookie_formats - with _visible_browser(self.config, publisher, viewport=None) as (context, page): + with _visible_browser( + self.config, + publisher, + viewport=None, + cancel_event=cancel_event, + ) as (context, page): def _try_save_captured() -> dict[str, Any] | None: """If a PDF was captured, save and validate it.""" + if _cancelled(cancel_event): + return None if captured_pdf: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(captured_pdf[-1]) - if is_pdf_file(output_path): + if write_pdf_bytes_atomic( + output_path, + captured_pdf[-1], + cancel_event, + ) and is_pdf_file(output_path): return _success(doi, output_path, "CARSI-Browser") return None @@ -187,8 +258,10 @@ def _try_save_captured() -> dict[str, Any] | None: pass # Capture PDF from network - captured_pdf = [] def on_response(response): + nonlocal capture_in_progress + if _cancelled(cancel_event): + return try: ct = response.headers.get("content-type", "") url = response.url @@ -198,21 +271,34 @@ def on_response(response): return if response.status >= 400: return - body = response.body() - if len(body) > 5000 and body[:4] == b"%PDF-": - captured_pdf.append(body) - log.info(f" [CARSI-Browser] PDF captured: {len(body)} bytes") + with capture_lock: + if _cancelled(cancel_event) or capture_in_progress or captured_pdf: + return + capture_in_progress = True + try: + body = response.body() + if _cancelled(cancel_event) or captured_pdf: + return + if len(body) > 5000 and body.startswith(b"%PDF-"): + captured_pdf.append(body) + log.info(f" [CARSI-Browser] PDF captured: {len(body)} bytes") + finally: + with capture_lock: + capture_in_progress = False except Exception: pass page.on("response", on_response) # Step 1: Navigate to article page first (gets Cloudflare clearance) log.info(f" [CARSI-Browser] Loading article: {article_url[:60]}") + if _cancelled(cancel_event): + return None try: page.goto(article_url, wait_until="domcontentloaded", timeout=60000) - time.sleep(5) except Exception: pass + if _wait_or_cancel(cancel_event, 5): + return None title = page.title() url = page.url @@ -221,9 +307,12 @@ def on_response(response): # Wait for Cloudflare challenge to resolve (visible stealth browser can pass it) from ..network import is_cloudflare_challenge for _cf_wait in range(12): + if _cancelled(cancel_event): + return None if is_cloudflare_challenge(page.title() or ""): log.info(f" [CARSI-Browser] Cloudflare challenge detected, waiting... ({_cf_wait+1}/12)") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None else: break else: @@ -281,21 +370,28 @@ def on_response(response): if not cookies_valid: # Step 2: Navigate to "Institutional login" link on article page + if _cancelled(cancel_event): + return None sso_href = page.evaluate(_SSO_LINK_FINDER_JS) if sso_href: log.info(f" [CARSI-Browser] Navigating to SSO: {sso_href[:80]}") + if _cancelled(cancel_event): + return None try: page.goto(sso_href, wait_until="domcontentloaded", timeout=30000) except Exception: pass else: log.info(" [CARSI-Browser] No SSO link found, trying direct login URL...") + if _cancelled(cancel_event): + return None try: page.goto(cfg.login_url, wait_until="domcontentloaded", timeout=30000) except Exception: pass - time.sleep(8) + if _wait_or_cancel(cancel_event, 8): + return None # Step 3: Search for institution in the WAYF page search_input = page.query_selector('#searchInstitution') @@ -307,17 +403,20 @@ def on_response(response): if search_input: search_input.fill(idp_en) - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None log.info(f" [CARSI-Browser] Searched for '{idp_en}'") # Click matching institution clicked = page.evaluate(_INSTITUTION_CLICK_JS, idp_en) if clicked: log.info(f" [CARSI-Browser] Selected: {clicked}") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None else: search_input.press("Enter") - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None else: log.info(" [CARSI-Browser] No institution search box found") @@ -330,7 +429,8 @@ def on_response(response): if any(x in url.lower() for x in _ak) or any(x in title for x in _at): log.info(" [CARSI-Browser] CAS login required. Please log in...") for i in range(100): - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None try: title = page.title() url = page.url @@ -353,13 +453,17 @@ def on_response(response): log.info(" [CARSI-Browser] Already authenticated") # Step 5: Navigate to article (with CARSI auth now) - time.sleep(2) + if _wait_or_cancel(cancel_event, 2): + return None log.info(f" [CARSI-Browser] Navigating to article: {article_url[:60]}") + if _cancelled(cancel_event): + return None try: page.goto(article_url, wait_until="domcontentloaded", timeout=30000) - time.sleep(5) except Exception: pass + if _wait_or_cancel(cancel_event, 5): + return None # Check for PDF via network capture saved = _try_save_captured() @@ -378,32 +482,44 @@ def on_response(response): if pdf_url and "{pii}" not in pdf_url: log.info(f" [CARSI-Browser] Trying PDF: {pdf_url[:80]}") captured_pdf.clear() + if _cancelled(cancel_event): + return None try: page.goto(pdf_url, wait_until="commit", timeout=30000) - time.sleep(5) except Exception: pass + if _wait_or_cancel(cancel_event, 5): + return None saved = _try_save_captured() if saved: return saved # Step 7: Find PDF link in HTML from ..pdf_utils import extract_pdf_url_from_html + if _cancelled(cancel_event): + return None html = page.content() + if _cancelled(cancel_event): + return None found_pdf = extract_pdf_url_from_html(html, page.url) if found_pdf: log.info(f" [CARSI-Browser] Found PDF link: {found_pdf[:80]}") captured_pdf.clear() + if _cancelled(cancel_event): + return None try: page.goto(found_pdf, wait_until="commit", timeout=30000) - time.sleep(5) except Exception: pass + if _wait_or_cancel(cancel_event, 5): + return None saved = _try_save_captured() if saved: return saved # Step 8: Click PDF button + if _cancelled(cancel_event): + return None click_result = page.evaluate(""" () => { const links = document.querySelectorAll('a'); @@ -422,17 +538,29 @@ def on_response(response): """) if click_result: log.info(f" [CARSI-Browser] Clicked: {str(click_result)[:80]}") - time.sleep(8) + if _wait_or_cancel(cancel_event, 8): + return None saved = _try_save_captured() if saved: return saved + if _cancelled(cancel_event): + return None log.info(f" [CARSI-Browser] No PDF found. Title: {page.title()[:40]} URL: {page.url[:60]}") return None except Exception as e: - log.info(f" [CARSI-Browser] Error: {e}") + if not _cancelled(cancel_event): + log.info(f" [CARSI-Browser] Error: {e}") return None + finally: + if page is not None and on_response is not None: + try: + page.remove_listener("response", on_response) + except Exception: + pass + with capture_lock: + captured_pdf.clear() def _find_downloaded_pdf(self, download_dir: str, doi: str) -> Path | None: """Check download directory for recently downloaded PDF files.""" @@ -486,6 +614,7 @@ def _validate_session(self, publisher: str) -> bool: # Validate by hitting a publisher page that requires auth # Use the main domain, not login_url (which always contains "login") + resp = None try: test_url = f"https://{cfg.domains[0]}/" resp = sess.get(test_url, timeout=15, allow_redirects=True) @@ -497,6 +626,12 @@ def _validate_session(self, publisher: str) -> bool: return resp.status_code == 200 except requests.RequestException: return False + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass def _browser_login(self, publisher: str) -> bool: """Login via CARSI using CloakBrowser.""" @@ -561,6 +696,9 @@ def _extract_chrome_cookies(self, publisher: str) -> None: log.warning(f" [CARSI] Chrome cookie extraction failed: {e}") def close(self): - for sess in self._sessions.values(): - sess.close() - self._sessions.clear() + try: + for sess in self._sessions.values(): + with contextlib.suppress(Exception): + sess.close() + finally: + self._sessions.clear() diff --git a/src/scansci_pdf/sources/carsi_source.py b/src/scansci_pdf/sources/carsi_source.py index dd9d634..c8fcfdb 100644 --- a/src/scansci_pdf/sources/carsi_source.py +++ b/src/scansci_pdf/sources/carsi_source.py @@ -5,6 +5,7 @@ """ from __future__ import annotations +import threading from pathlib import Path from typing import Any @@ -13,11 +14,18 @@ log = get_logger() -def try_carsi(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: +def try_carsi( + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Try downloading via CARSI federated auth without WebVPN dependency. Returns a result dict on success, None on failure. """ + if cancel_event is not None and cancel_event.is_set(): + return None if not config.get("carsi_enabled", False): return None @@ -25,11 +33,14 @@ def try_carsi(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, if not idp_name: return None + client = None try: from .carsi import CARSIClient, detect_publisher from .instsci import _resolve_doi_url resolved_url = _resolve_doi_url(doi) + if cancel_event is not None and cancel_event.is_set(): + return None if not resolved_url: resolved_url = f"https://doi.org/{doi}" @@ -55,11 +66,22 @@ def try_carsi(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, log.info(f" [CARSI] Redirected to primary domain: {resolved_url[:80]}") # Try browser download (CloakBrowser first, Selenium fallback) - result = client.download_via_browser(doi, resolved_url, output_path) + if cancel_event is not None and cancel_event.is_set(): + return None + result = client.download_via_browser( + doi, + resolved_url, + output_path, + cancel_event=cancel_event, + ) if result: return result except ImportError: return None except Exception as e: - log.info(f" [CARSI] {e}") + if cancel_event is None or not cancel_event.is_set(): + log.info(f" [CARSI] {e}") + finally: + if client is not None: + client.close() return None diff --git a/src/scansci_pdf/sources/ezproxy.py b/src/scansci_pdf/sources/ezproxy.py index 420af04..e0f7a9f 100644 --- a/src/scansci_pdf/sources/ezproxy.py +++ b/src/scansci_pdf/sources/ezproxy.py @@ -7,6 +7,7 @@ from __future__ import annotations +import threading import time from pathlib import Path from typing import Any @@ -20,11 +21,38 @@ is_pdf_file, is_plausible_pdf_url, success, + write_pdf_bytes_atomic, ) log = get_logger() +def _cancelled(cancel_event: threading.Event | None) -> bool: + return cancel_event is not None and cancel_event.is_set() + + +def _wait_or_cancel( + cancel_event: threading.Event | None, + timeout: float, +) -> bool: + if cancel_event is None: + time.sleep(max(0.0, timeout)) + return False + return cancel_event.wait(max(0.0, timeout)) + + +def _close_browser_resource(resource: Any) -> None: + """Best-effort close on the resource's owner thread, with one retry.""" + if resource is None: + return + for _attempt in range(2): + try: + resource.close() + return + except Exception: + pass + + def _get_ezproxy_base(config: dict[str, Any]) -> str: """Get EZProxy login URL template.""" return config.get("ezproxy_login_url", "") @@ -51,16 +79,23 @@ def _validate_ezproxy_session(config: dict[str, Any]) -> bool: if not cookies: return False - sess = requests.Session() - sess.trust_env = False - for c in cookies: - sess.cookies.set(c["name"], c["value"], domain=c.get("domain", ""), path=c.get("path", "/")) - - # Test with a known URL - test_url = _make_ezproxy_url("https://www.sciencedirect.com", config) - if not test_url: - return False + sess = None + resp = None try: + sess = requests.Session() + sess.trust_env = False + for c in cookies: + sess.cookies.set( + c["name"], + c["value"], + domain=c.get("domain", ""), + path=c.get("path", "/"), + ) + + # Test with a known URL + test_url = _make_ezproxy_url("https://www.sciencedirect.com", config) + if not test_url: + return False resp = sess.get(test_url, timeout=15, allow_redirects=True) # If redirected to login, session is invalid if "login" in resp.url.lower() or "libproxy" in resp.url.lower(): @@ -68,6 +103,17 @@ def _validate_ezproxy_session(config: dict[str, Any]) -> bool: return resp.status_code == 200 except Exception: return False + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass + if sess is not None: + try: + sess.close() + except Exception: + pass def _ezproxy_cookie_path(config: dict[str, Any]) -> Path: @@ -91,13 +137,18 @@ def ezproxy_login(config: dict[str, Any]) -> bool: return False -def try_ezproxy(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: +def try_ezproxy( + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Try downloading paper through EZProxy institutional proxy. Uses Selenium browser to access the paper through the library proxy, which handles authentication and cookie management automatically. """ - if not config.get("ezproxy_enabled", False): + if _cancelled(cancel_event) or not config.get("ezproxy_enabled", False): return None base = _get_ezproxy_base(config) @@ -105,11 +156,22 @@ def try_ezproxy(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str return None # Resolve DOI to get publisher URL + resp = None try: + if _cancelled(cancel_event): + return None resp = requests.head(f"https://doi.org/{doi}", allow_redirects=True, timeout=10) resolved_url = resp.url except Exception: resolved_url = f"https://doi.org/{doi}" + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass + if _cancelled(cancel_event): + return None # Construct EZProxy URL ezproxy_url = _make_ezproxy_url(resolved_url, config) @@ -119,33 +181,73 @@ def try_ezproxy(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str log.info(f" [EZProxy] Trying {doi} via library proxy...") try: + from ..cloakbrowser_compat import ( + launch_with_driver_cleanup, + prepare_cloakbrowser_runtime, + ) + + prepare_cloakbrowser_runtime() from cloakbrowser import launch from ..browser_engine import _build_browser_args - except ImportError: + except Exception: log.info(" [EZProxy] cloakbrowser not installed") return None download_dir = str(output_path.parent) args = _build_browser_args(config) captured_pdf: list[bytes] = [] + capture_lock = threading.Lock() + capture_in_progress = False def _on_response(response): + nonlocal capture_in_progress + if _cancelled(cancel_event): + return try: ct = response.headers.get("content-type", "") if "pdf" in ct: try: - body = response.body() - if body and len(body) > 5000: - captured_pdf.append(body) + with capture_lock: + if _cancelled(cancel_event) or capture_in_progress or captured_pdf: + return + capture_in_progress = True + try: + body = response.body() + if not _cancelled(cancel_event) and not captured_pdf and body and len(body) > 5000: + captured_pdf.append(body) + finally: + with capture_lock: + capture_in_progress = False except Exception: pass except Exception: pass - browser = launch(headless=False, humanize=True, args=args) + browser = None + context = None + page = None + slot_lease = None try: + if _cancelled(cancel_event): + return None + from .. import browser_engine + slot_lease = browser_engine._retain_browser_slot(config, cancel_event) + raw_browser = launch_with_driver_cleanup( + launch, + headless=False, + humanize=True, + args=args, + ) + browser = browser_engine._LeasedBrowser(raw_browser, slot_lease) + slot_lease = None + if _cancelled(cancel_event): + return None context = browser.new_context() + if _cancelled(cancel_event): + return None page = context.new_page() + if _cancelled(cancel_event): + return None page.on("response", _on_response) # Load saved cookies if available @@ -157,8 +259,11 @@ def _on_response(response): context.add_cookies(cookies) # Navigate to EZProxy URL + if _cancelled(cancel_event): + return None page.goto(ezproxy_url, wait_until="domcontentloaded", timeout=30000) - time.sleep(8) + if _wait_or_cancel(cancel_event, 8): + return None # Check if redirected to login url = page.url @@ -167,7 +272,8 @@ def _on_response(response): max_wait = 180 elapsed = 0 while elapsed < max_wait: - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None elapsed += 3 try: url = page.url @@ -180,14 +286,15 @@ def _on_response(response): return None # Check for captured PDF - if captured_pdf: + if not _cancelled(cancel_event) and captured_pdf: pdf_bytes = captured_pdf[-1] if pdf_bytes[:5] == b"%PDF-": - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - return success(doi, output_path, "EZProxy") + if write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + return success(doi, output_path, "EZProxy") # Look for PDF link in page + if _cancelled(cancel_event): + return None pdf_link = page.evaluate("""() => { const links = document.querySelectorAll('a'); for (const link of links) { @@ -199,25 +306,37 @@ def _on_response(response): } return ''; }""") + if _cancelled(cancel_event): + return None if pdf_link: log.info(f" [EZProxy] Found PDF link: {pdf_link[:80]}") captured_pdf.clear() + if _cancelled(cancel_event): + return None page.goto(pdf_link, wait_until="commit", timeout=30000) - time.sleep(5) - if captured_pdf: + if _wait_or_cancel(cancel_event, 5): + return None + if not _cancelled(cancel_event) and captured_pdf: pdf_bytes = captured_pdf[-1] if pdf_bytes[:5] == b"%PDF-": - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - return success(doi, output_path, "EZProxy") + if write_pdf_bytes_atomic(output_path, pdf_bytes, cancel_event): + return success(doi, output_path, "EZProxy") except Exception as e: - log.info(f" [EZProxy] Error: {e}") + if not _cancelled(cancel_event): + log.info(f" [EZProxy] Error: {e}") finally: - try: - browser.close() - except Exception: - pass + if page is not None: + try: + page.remove_listener("response", _on_response) + except Exception: + pass + with capture_lock: + captured_pdf.clear() + for resource in (page, context, browser): + _close_browser_resource(resource) + if browser is None and slot_lease is not None: + slot_lease.close() return None diff --git a/src/scansci_pdf/sources/instsci.py b/src/scansci_pdf/sources/instsci.py index c2ebbbe..2ff8840 100644 --- a/src/scansci_pdf/sources/instsci.py +++ b/src/scansci_pdf/sources/instsci.py @@ -13,6 +13,7 @@ import binascii import json import re +import threading import time import urllib.parse from pathlib import Path @@ -23,11 +24,14 @@ from ..log import get_logger from ..pdf_utils import ( + _bind_session_to_response, _response_looks_pdf, extract_pdf_url_from_html, is_pdf_file, is_plausible_pdf_url, success, + write_pdf_bytes_atomic, + write_pdf_stream_atomic, ) # Import compiled core functions if available (Cython .pyd/.so) @@ -59,15 +63,45 @@ def _cfg(config: dict[str, Any], suffix: str, default: Any = None) -> Any: _INSTSCI_DELAY_MAX = 5.0 +def _cancelled(cancel_event: threading.Event | None) -> bool: + return cancel_event is not None and cancel_event.is_set() -def _instsci_rate_limit() -> None: + +def _wait_or_cancel( + cancel_event: threading.Event | None, + timeout: float, +) -> bool: + """Wait for ``timeout`` seconds; return True when cancellation wins.""" + if cancel_event is None: + time.sleep(max(0.0, timeout)) + return False + return cancel_event.wait(max(0.0, timeout)) + + +def _close_browser_resource(resource: Any) -> None: + """Best-effort close on the resource's owner thread, with one retry.""" + if resource is None: + return + for _attempt in range(2): + try: + resource.close() + return + except Exception: + pass + + + +def _instsci_rate_limit(cancel_event: threading.Event | None = None) -> bool: global _last_instsci_time + if _cancelled(cancel_event): + return False now = time.time() elapsed = now - _last_instsci_time delay = __import__("random").uniform(_INSTSCI_DELAY_MIN, _INSTSCI_DELAY_MAX) - if elapsed < delay: - time.sleep(delay - elapsed) + if elapsed < delay and _wait_or_cancel(cancel_event, delay - elapsed): + return False _last_instsci_time = time.time() + return True def instsci_cookie_path(config: dict[str, Any]) -> Path: @@ -203,6 +237,8 @@ def _validate_session(config: dict[str, Any]) -> bool: if not base: return False test_url = convert_url("https://www.nature.com", base, config) + s = None + resp = None try: s = requests.Session() s.trust_env = False @@ -214,6 +250,17 @@ def _validate_session(config: dict[str, Any]) -> bool: return resp.status_code == 200 except Exception: return False + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass + if s is not None: + try: + s.close() + except Exception: + pass def instsci_login(config: dict[str, Any]) -> bool: @@ -274,7 +321,17 @@ def _fetch_via_webvpn(url: str, config: dict[str, Any], *, stream: bool = False) s.proxies = {"http": socks5, "https": socks5} s.verify = False # Campus connectors often use self-signed certs - return s.get(proxied, timeout=request_timeout(config), allow_redirects=True, stream=stream) + try: + resp = s.get( + proxied, + timeout=request_timeout(config), + allow_redirects=True, + stream=stream, + ) + except Exception: + s.close() + raise + return _bind_session_to_response(resp, s) def _fetch_direct_via_socks5(url: str, config: dict[str, Any], *, stream: bool = False) -> requests.Response: @@ -288,11 +345,22 @@ def _fetch_direct_via_socks5(url: str, config: dict[str, Any], *, stream: bool = s.proxies = {"http": socks5, "https": socks5} s.verify = False - return s.get(url, timeout=request_timeout(config), allow_redirects=True, stream=stream) + try: + resp = s.get( + url, + timeout=request_timeout(config), + allow_redirects=True, + stream=stream, + ) + except Exception: + s.close() + raise + return _bind_session_to_response(resp, s) def _resolve_doi_url(doi: str) -> str | None: """Resolve DOI to get the publisher URL.""" + resp = None try: resp = requests.get( f"https://doi.org/{doi}", @@ -302,11 +370,17 @@ def _resolve_doi_url(doi: str) -> str | None: stream=True, verify=False, ) - resp.close() - if resp.url and resp.url != f"https://doi.org/{doi}": - return resp.url + final_url = resp.url + if final_url and final_url != f"https://doi.org/{doi}": + return final_url except Exception: pass + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass return None @@ -522,15 +596,22 @@ def _extract_inline_pdf(page: Any) -> bytes | None: def _download_pdf_with_browser_cookies( - pdf_url: str, output_path: Path, config: dict[str, Any], doi: str, context: Any + pdf_url: str, + output_path: Path, + config: dict[str, Any], + doi: str, + context: Any, + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Download PDF via WebVPN using cookies from a live browser context. Bypasses stale cookie files by pulling cookies directly from the Playwright browser context that was just used for login/page-navigation. """ - if not is_plausible_pdf_url(pdf_url): + if _cancelled(cancel_event) or not is_plausible_pdf_url(pdf_url): return None + session = None + resp = None try: from ..network import USER_AGENT, request_timeout @@ -539,59 +620,85 @@ def _download_pdf_with_browser_cookies( return None proxied = convert_url(pdf_url, base, config) - s = requests.Session() - s.trust_env = False - s.headers.update({"User-Agent": USER_AGENT}) + session = requests.Session() + session.trust_env = False + session.headers.update({"User-Agent": USER_AGENT}) for c in context.cookies(): + if _cancelled(cancel_event): + return None name = c.get("name") value = c.get("value") if name and value is not None: - s.cookies.set(name, value, domain=c.get("domain", ""), path=c.get("path", "/")) + session.cookies.set(name, value, domain=c.get("domain", ""), path=c.get("path", "/")) - resp = s.get(proxied, timeout=request_timeout(config), allow_redirects=True, stream=True) - if resp.status_code >= 400: + if _cancelled(cancel_event): + return None + resp = session.get( + proxied, + timeout=request_timeout(config), + allow_redirects=True, + stream=True, + ) + if _cancelled(cancel_event) or resp.status_code >= 400: log.info(f" [WebVPN-Browser] Browser-cookie HTTP status={resp.status_code} for PDF URL") return None iterator = resp.iter_content(chunk_size=8192) first = next(iterator, b"") + if _cancelled(cancel_event): + return None if not _response_looks_pdf(resp, first): return None - output_path.parent.mkdir(parents=True, exist_ok=True) - tmp = output_path.with_suffix(output_path.suffix + ".part") - try: - with tmp.open("wb") as fh: - fh.write(first) - for chunk in iterator: - if chunk: - fh.write(chunk) - tmp.replace(output_path) - except Exception: - tmp.unlink(missing_ok=True) - raise + if not write_pdf_stream_atomic( + output_path, + first, + iterator, + cancel_event, + ): + return None - if is_pdf_file(output_path): + if not _cancelled(cancel_event) and is_pdf_file(output_path): log.info(f" [WebVPN-Browser] PDF downloaded via browser-cookie HTTP") return success(doi, output_path, "WebVPN(Browser)") except Exception as e: - log.info(f" [WebVPN-Browser] Browser-cookie HTTP error: {e}") + if not _cancelled(cancel_event): + log.info(f" [WebVPN-Browser] Browser-cookie HTTP error: {e}") + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass + if session is not None: + try: + session.close() + except Exception: + pass return None -def _try_instsci_socks5(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: +def _try_instsci_socks5( + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Try downloading via SOCKS5 campus connector (EasyConnect/aTrust). The SOCKS5 connector handles authentication at the network level, so no WebVPN URL conversion or CAS login is needed — just fetch the publisher URL directly through the proxy. """ - _instsci_rate_limit() + if not _instsci_rate_limit(cancel_event): + return None socks5 = _get_socks5_proxy(config) log.info(f" [CampusConnector] Trying {doi} via {socks5}") # Step 1: Resolve DOI to publisher URL resolved_url = _resolve_doi_url(doi) + if _cancelled(cancel_event): + return None if not resolved_url: resolved_url = f"https://doi.org/{doi}" @@ -599,57 +706,86 @@ def _try_instsci_socks5(doi: str, output_path: Path, config: dict[str, Any]) -> pdf_url = _construct_publisher_pdf_url(doi, resolved_url) if pdf_url: log.info(f" [CampusConnector] Trying publisher PDF: {pdf_url[:80]}") - result = _download_pdf_socks5(pdf_url, output_path, config, doi) + result = _download_pdf_socks5( + pdf_url, + output_path, + config, + doi, + cancel_event=cancel_event, + ) if result: return result # Step 3: Fetch landing page via SOCKS5 and find PDF link + resp = None try: + if _cancelled(cancel_event): + return None resp = _fetch_direct_via_socks5(resolved_url, config, stream=True) + if _cancelled(cancel_event): + return None if resp.status_code >= 400: log.info(f" [CampusConnector] HTTP {resp.status_code} for {resolved_url[:60]}") return None iterator = resp.iter_content(chunk_size=8192) first = next(iterator, b"") + if _cancelled(cancel_event): + return None # Direct PDF response if _response_looks_pdf(resp, first): - output_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = output_path.with_suffix(output_path.suffix + ".part") - try: - with tmp_path.open("wb") as fh: - fh.write(first) - for chunk in iterator: - if chunk: - fh.write(chunk) - tmp_path.replace(output_path) - except Exception: - tmp_path.unlink(missing_ok=True) - raise - if is_pdf_file(output_path): + if not write_pdf_stream_atomic( + output_path, + first, + iterator, + cancel_event, + ): + return None + if not _cancelled(cancel_event) and is_pdf_file(output_path): log.info(" [CampusConnector] PDF downloaded directly") return success(doi, output_path, "CampusConnector") # HTML response - look for PDF link html = first + resp.raw.read(512_000, decode_content=True) + if _cancelled(cancel_event): + return None html_str = html.decode("utf-8", errors="ignore") found_pdf = _find_pdf_link(html_str, resp.url) if found_pdf: log.info(f" [CampusConnector] Found PDF link: {found_pdf[:80]}") - result = _download_pdf_socks5(found_pdf, output_path, config, doi) + result = _download_pdf_socks5( + found_pdf, + output_path, + config, + doi, + cancel_event=cancel_event, + ) if result: return result pdf_url = extract_pdf_url_from_html(html_str, resp.url) if pdf_url: - result = _download_pdf_socks5(pdf_url, output_path, config, doi) + result = _download_pdf_socks5( + pdf_url, + output_path, + config, + doi, + cancel_event=cancel_event, + ) if result: return result except Exception as e: - log.info(f" [CampusConnector] {e}") + if not _cancelled(cancel_event): + log.info(f" [CampusConnector] {e}") + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass return None @@ -659,47 +795,67 @@ def _download_pdf_socks5( output_path: Path, config: dict[str, Any], doi: str, + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: """Download a PDF URL directly through SOCKS5 campus connector.""" - if not is_plausible_pdf_url(url): + if _cancelled(cancel_event) or not is_plausible_pdf_url(url): return None + resp = None try: - _instsci_rate_limit() + if not _instsci_rate_limit(cancel_event): + return None resp = _fetch_direct_via_socks5(url, config, stream=True) - if resp.status_code >= 400: + if _cancelled(cancel_event) or resp.status_code >= 400: return None iterator = resp.iter_content(chunk_size=8192) first_chunk = next(iterator, b"") + if _cancelled(cancel_event): + return None if not _response_looks_pdf(resp, first_chunk): return None - output_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = output_path.with_suffix(output_path.suffix + ".part") - try: - with tmp_path.open("wb") as fh: - fh.write(first_chunk) - for chunk in iterator: - if chunk: - fh.write(chunk) - tmp_path.replace(output_path) - except Exception: - tmp_path.unlink(missing_ok=True) - raise + if not write_pdf_stream_atomic( + output_path, + first_chunk, + iterator, + cancel_event, + ): + return None - if is_pdf_file(output_path): + if not _cancelled(cancel_event) and is_pdf_file(output_path): log.info(f" [CampusConnector] PDF downloaded: {doi}") return success(doi, output_path, "CampusConnector") except Exception as e: - log.info(f" [CampusConnector] Download error: {e}") + if not _cancelled(cancel_event): + log.info(f" [CampusConnector] Download error: {e}") + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass return None -def _try_instsci_browser(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: +def _try_instsci_browser( + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Download via visible stealth browser browser. Login + download in same session.""" + if _cancelled(cancel_event): + return None try: + from ..cloakbrowser_compat import ( + launch_with_driver_cleanup, + prepare_cloakbrowser_runtime, + ) + + prepare_cloakbrowser_runtime() from cloakbrowser import launch - except ImportError: + except Exception: log.info(" [WebVPN-Browser] cloakbrowser not installed") return None @@ -708,6 +864,8 @@ def _try_instsci_browser(doi: str, output_path: Path, config: dict[str, Any]) -> return None resolved_url = _resolve_doi_url(doi) + if _cancelled(cancel_event): + return None if not resolved_url: resolved_url = f"https://doi.org/{doi}" @@ -716,11 +874,35 @@ def _try_instsci_browser(doi: str, output_path: Path, config: dict[str, Any]) -> print(f"\n [WebVPN] 正在打开浏览器，请在浏览器中登录 WebVPN...") print(f" 登录完成后等待 5 秒，程序会自动继续下载。\n") + browser = None + context = None + page = None + slot_lease = None + captured_pdf: list[bytes] = [] + capture_lock = threading.Lock() + capture_in_progress = False + on_response = None try: - browser = launch(headless=False, humanize=True, - args=["--disable-features=CrossOriginOpenerPolicy"]) + if _cancelled(cancel_event): + return None + from .. import browser_engine + slot_lease = browser_engine._retain_browser_slot(config, cancel_event) + raw_browser = launch_with_driver_cleanup( + launch, + headless=False, + humanize=True, + args=["--disable-features=CrossOriginOpenerPolicy"], + ) + browser = browser_engine._LeasedBrowser(raw_browser, slot_lease) + slot_lease = None + if _cancelled(cancel_event): + return None context = browser.new_context() + if _cancelled(cancel_event): + return None page = context.new_page() + if _cancelled(cancel_event): + return None # Restore saved cookies before navigating cookie_path = instsci_cookie_path(config) @@ -750,9 +932,10 @@ def _try_instsci_browser(doi: str, output_path: Path, config: dict[str, Any]) -> pass # Capture PDF from network responses - captured_pdf = [] - def on_response(response): + nonlocal capture_in_progress + if _cancelled(cancel_event): + return try: ct = response.headers.get("content-type", "") url = response.url @@ -763,10 +946,20 @@ def on_response(response): return if response.status >= 400: return - body = response.body() - if len(body) > 5000 and body[:4] == b"%PDF-": - captured_pdf.append(body) - log.info(f" [WebVPN-Browser] PDF captured: {len(body)} bytes from {url[:60]}") + with capture_lock: + if _cancelled(cancel_event) or capture_in_progress or captured_pdf: + return + capture_in_progress = True + try: + body = response.body() + if _cancelled(cancel_event) or captured_pdf: + return + if len(body) > 5000 and body.startswith(b"%PDF-"): + captured_pdf.append(body) + log.info(f" [WebVPN-Browser] PDF captured: {len(body)} bytes from {url[:60]}") + finally: + with capture_lock: + capture_in_progress = False except Exception: pass @@ -774,11 +967,14 @@ def on_response(response): # Navigate to paper URL directly via WebVPN # If not logged in, will redirect to login page + if _cancelled(cancel_event): + return None try: page.goto(webvpn_url, wait_until="domcontentloaded", timeout=60000) except Exception: pass - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None # If on login page, wait for user to login then retry title = page.title() @@ -792,7 +988,8 @@ def on_response(response): print(f" 检测到登录页面，请完成登录...") # Wait up to 5 minutes, checking title every 3 seconds for i in range(100): - time.sleep(3) + if _wait_or_cancel(cancel_event, 3): + return None try: title = page.title() url_now = page.url @@ -837,19 +1034,33 @@ def on_response(response): # Try HTTP download with fresh browser cookies (bypasses JS-heavy pages) pdf_url = _construct_publisher_pdf_url(doi, resolved_url) if pdf_url: - time.sleep(2) - result = _download_pdf_with_browser_cookies(pdf_url, output_path, config, doi, context) + if _wait_or_cancel(cancel_event, 2): + return None + result = _download_pdf_with_browser_cookies( + pdf_url, + output_path, + config, + doi, + context, + cancel_event=cancel_event, + ) if result: return result # Fall back to browser-based PDF extraction - time.sleep(2) + if _wait_or_cancel(cancel_event, 2): + return None + if _cancelled(cancel_event): + return None try: page.goto(webvpn_url, wait_until="domcontentloaded", timeout=60000) except Exception: pass + if _cancelled(cancel_event): + return None for _w in range(20): - time.sleep(2) + if _wait_or_cancel(cancel_event, 2): + return None try: _t = (page.title() or "").lower() _b = (page.evaluate("document.body?.innerText?.length || 0") or 0) @@ -861,22 +1072,37 @@ def on_response(response): # Already authenticated — try HTTP download with browser cookies pdf_url = _construct_publisher_pdf_url(doi, resolved_url) if pdf_url: - time.sleep(2) - result = _download_pdf_with_browser_cookies(pdf_url, output_path, config, doi, context) + if _wait_or_cancel(cancel_event, 2): + return None + result = _download_pdf_with_browser_cookies( + pdf_url, + output_path, + config, + doi, + context, + cancel_event=cancel_event, + ) if result: return result - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None # Helper: try to save captured PDF def _save_captured(): + if _cancelled(cancel_event): + return None if captured_pdf: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(captured_pdf[-1]) - if is_pdf_file(output_path): + if write_pdf_bytes_atomic( + output_path, + captured_pdf[-1], + cancel_event, + ) and is_pdf_file(output_path): return success(doi, output_path, "WebVPN(Browser)") return None # Check if page itself is a PDF (inline viewer) + if _cancelled(cancel_event): + return None page_url = page.url page_title = page.title() log.info(f" [WebVPN-Browser] On page: title='{page_title[:40]}' url={page_url[:60]}") @@ -889,17 +1115,30 @@ def _save_captured(): # If page looks like inline PDF viewer, try to get the PDF bytes if _is_inline_pdf_page(page): pdf_bytes = _extract_inline_pdf(page) + if _cancelled(cancel_event): + return None if pdf_bytes: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - if is_pdf_file(output_path): + if write_pdf_bytes_atomic( + output_path, + pdf_bytes, + cancel_event, + ) and is_pdf_file(output_path): return success(doi, output_path, "WebVPN(Browser)") # Strategy 1: Try direct publisher PDF URL via browser-cookie HTTP first resolved_for_pdf = _resolve_doi_url(doi) or f"https://doi.org/{doi}" + if _cancelled(cancel_event): + return None pdf_url = _construct_publisher_pdf_url(doi, resolved_for_pdf) if pdf_url: - result = _download_pdf_with_browser_cookies(pdf_url, output_path, config, doi, context) + result = _download_pdf_with_browser_cookies( + pdf_url, + output_path, + config, + doi, + context, + cancel_event=cancel_event, + ) if result: return result @@ -907,66 +1146,101 @@ def _save_captured(): pdf_webvpn = convert_url(pdf_url, base, config) log.info(f" [WebVPN-Browser] Trying direct PDF via browser: {pdf_webvpn[:80]}") captured_pdf.clear() + if _cancelled(cancel_event): + return None try: with page.expect_download(timeout=30000) as download_info: page.goto(pdf_webvpn, wait_until="commit", timeout=30000) + if _cancelled(cancel_event): + return None download = download_info.value tmp = download.path() pdf_bytes = tmp.read_bytes() if tmp else None - if pdf_bytes and pdf_bytes[:4] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - if is_pdf_file(output_path): + if pdf_bytes and pdf_bytes.startswith(b"%PDF-") and len(pdf_bytes) > 5000: + if write_pdf_bytes_atomic( + output_path, + pdf_bytes, + cancel_event, + ) and is_pdf_file(output_path): return success(doi, output_path, "WebVPN(Browser)") except Exception as dl_exc: log.info(f" [WebVPN-Browser] Download event not triggered: {dl_exc}") - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None result = _save_captured() if result: return result if _is_inline_pdf_page(page): pdf_bytes = _extract_inline_pdf(page) + if _cancelled(cancel_event): + return None if pdf_bytes: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - if is_pdf_file(output_path): + if write_pdf_bytes_atomic( + output_path, + pdf_bytes, + cancel_event, + ) and is_pdf_file(output_path): return success(doi, output_path, "WebVPN(Browser)") # Strategy 2: Find PDF link in HTML, try browser-cookie HTTP first + if _cancelled(cancel_event): + return None html = page.content() + if _cancelled(cancel_event): + return None found_pdf_url = extract_pdf_url_from_html(html, page.url) if found_pdf_url: log.info(f" [WebVPN-Browser] Found PDF link: {found_pdf_url[:80]}") - result = _download_pdf_with_browser_cookies(found_pdf_url, output_path, config, doi, context) + result = _download_pdf_with_browser_cookies( + found_pdf_url, + output_path, + config, + doi, + context, + cancel_event=cancel_event, + ) if result: return result # Fallback: expect_download in browser captured_pdf.clear() + if _cancelled(cancel_event): + return None try: with page.expect_download(timeout=30000) as download_info: page.goto(found_pdf_url, wait_until="commit", timeout=30000) + if _cancelled(cancel_event): + return None download = download_info.value tmp = download.path() pdf_bytes = tmp.read_bytes() if tmp else None - if pdf_bytes and pdf_bytes[:4] == b"%PDF-" and len(pdf_bytes) > 5000: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - if is_pdf_file(output_path): + if pdf_bytes and pdf_bytes.startswith(b"%PDF-") and len(pdf_bytes) > 5000: + if write_pdf_bytes_atomic( + output_path, + pdf_bytes, + cancel_event, + ) and is_pdf_file(output_path): return success(doi, output_path, "WebVPN(Browser)") except Exception: - time.sleep(5) + if _wait_or_cancel(cancel_event, 5): + return None result = _save_captured() if result: return result if _is_inline_pdf_page(page): pdf_bytes = _extract_inline_pdf(page) + if _cancelled(cancel_event): + return None if pdf_bytes: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(pdf_bytes) - if is_pdf_file(output_path): + if write_pdf_bytes_atomic( + output_path, + pdf_bytes, + cancel_event, + ) and is_pdf_file(output_path): return success(doi, output_path, "WebVPN(Browser)") + if _cancelled(cancel_event): + return None log.info(f" [WebVPN-Browser] No PDF found. Title: {page.title()[:40]} URL: {page.url[:60]}") return None @@ -974,13 +1248,25 @@ def _save_captured(): log.info(f" [WebVPN-Browser] Error: {e}") return None finally: - try: - browser.close() - except Exception: - pass + if page is not None and on_response is not None: + try: + page.remove_listener("response", on_response) + except Exception: + pass + with capture_lock: + captured_pdf.clear() + for resource in (page, context, browser): + _close_browser_resource(resource) + if browser is None and slot_lease is not None: + slot_lease.close() -def try_instsci(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: +def try_instsci( + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Try downloading paper through institutional access. Strategy: @@ -991,17 +1277,29 @@ def try_instsci(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str Note: CARSI is now a standalone source tier (carsi_source.try_carsi), called independently from the download orchestrator. """ - if not _cfg(config, "enabled", False): + if _cancelled(cancel_event) or not _cfg(config, "enabled", False): return None # Step 0: SOCKS5 campus connector mode (EasyConnect/aTrust) — direct access if _is_campus_connector_mode(config): - result = _try_instsci_socks5(doi, output_path, config) + result = _try_instsci_socks5( + doi, + output_path, + config, + cancel_event=cancel_event, + ) if result: return result # Step 1: Try stealth browser download (handles CAS auth + Cloudflare) - result = _try_instsci_browser(doi, output_path, config) + if _cancelled(cancel_event): + return None + result = _try_instsci_browser( + doi, + output_path, + config, + cancel_event=cancel_event, + ) if result: return result @@ -1009,8 +1307,13 @@ def try_instsci(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str # _validate_session fails — the stealth browser may have just logged in # and saved fresh cookies that work for the target paper but fail # validation's unrelated test URL) - if instsci_cookie_path(config).exists(): - result = _try_instsci_http(doi, output_path, config) + if not _cancelled(cancel_event) and instsci_cookie_path(config).exists(): + result = _try_instsci_http( + doi, + output_path, + config, + cancel_event=cancel_event, + ) if result: return result @@ -1018,15 +1321,23 @@ def try_instsci(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str return None -def _try_instsci_http(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: +def _try_instsci_http( + doi: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Try downloading via HTTP with saved cookies.""" - _instsci_rate_limit() + if not _instsci_rate_limit(cancel_event): + return None log.info(f" [WebVPN] Trying {doi}") # Step 1: Resolve DOI to get publisher URL resolved_url = _resolve_doi_url(doi) + if _cancelled(cancel_event): + return None if not resolved_url: resolved_url = f"https://doi.org/{doi}" @@ -1034,46 +1345,58 @@ def _try_instsci_http(doi: str, output_path: Path, config: dict[str, Any]) -> di pdf_url = _construct_publisher_pdf_url(doi, resolved_url) if pdf_url: log.info(f" [WebVPN] Trying publisher PDF: {pdf_url[:80]}...") - result = _download_pdf_instsci(pdf_url, output_path, config, doi) + result = _download_pdf_instsci( + pdf_url, + output_path, + config, + doi, + cancel_event=cancel_event, + ) if result: return result # Step 3: Fetch via WebVPN and look for PDF link in HTML + resp = None try: doi_url = f"https://doi.org/{doi}" + if _cancelled(cancel_event): + return None resp = _fetch_via_webvpn(doi_url, config, stream=True) - if resp.status_code >= 400: + if _cancelled(cancel_event) or resp.status_code >= 400: return None iterator = resp.iter_content(chunk_size=8192) first = next(iterator, b"") + if _cancelled(cancel_event): + return None # Direct PDF response if _response_looks_pdf(resp, first): - output_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = output_path.with_suffix(output_path.suffix + ".part") - try: - with tmp_path.open("wb") as fh: - fh.write(first) - for chunk in iterator: - if chunk: - fh.write(chunk) - tmp_path.replace(output_path) - except Exception: - tmp_path.unlink(missing_ok=True) - raise - if is_pdf_file(output_path): + if not write_pdf_stream_atomic( + output_path, + first, + iterator, + cancel_event, + ): + return None + if not _cancelled(cancel_event) and is_pdf_file(output_path): return success(doi, output_path, "WebVPN") # HTML response - extract PDF link html = first + resp.raw.read(512_000, decode_content=True) + if _cancelled(cancel_event): + return None html_str = html.decode("utf-8", errors="ignore") # Check for Cloudflare block from ..network import _is_cloudflare_block if any(sig in html_str.lower() for sig in ("cf-browser-verification", "challenge-platform", "just a moment", "请稍候", "正在验证", "checking your browser")): log.info(" [WebVPN] Cloudflare detected, trying browser...") - browser_html = _try_browser_via_webvpn(doi_url, config) + browser_html = _try_browser_via_webvpn( + doi_url, + config, + cancel_event=cancel_event, + ) if browser_html: html_str = browser_html @@ -1081,25 +1404,51 @@ def _try_instsci_http(doi: str, output_path: Path, config: dict[str, Any]) -> di found_pdf = _find_pdf_link(html_str, resp.url) if found_pdf: log.info(f" [WebVPN] Found PDF link in HTML: {found_pdf[:80]}...") - result = _download_pdf_instsci(found_pdf, output_path, config, doi) + result = _download_pdf_instsci( + found_pdf, + output_path, + config, + doi, + cancel_event=cancel_event, + ) if result: return result # Fallback to extract_pdf_url_from_html pdf_url = extract_pdf_url_from_html(html_str, resp.url) if pdf_url: - return _download_pdf_instsci(pdf_url, output_path, config, doi) + return _download_pdf_instsci( + pdf_url, + output_path, + config, + doi, + cancel_event=cancel_event, + ) except Exception as e: - log.info(f" [WebVPN] {e}") + if not _cancelled(cancel_event): + log.info(f" [WebVPN] {e}") + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass return None -def _try_carsi(doi: str, resolved_url: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: +def _try_carsi( + doi: str, + resolved_url: str, + output_path: Path, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Try downloading via CARSI federated auth (browser-based).""" - if not config.get("carsi_enabled", False): + if _cancelled(cancel_event) or not config.get("carsi_enabled", False): return None + client = None try: from .carsi import CARSIClient, detect_publisher publisher = detect_publisher(resolved_url) @@ -1109,55 +1458,71 @@ def _try_carsi(doi: str, resolved_url: str, output_path: Path, config: dict[str, # Try stealth browser first (stealth browser, handles Cloudflare) log.info(f" [CARSI] Trying browser download for {doi}...") - result = client.download_via_browser(doi, resolved_url, output_path) - if result: - return result - - # Fallback to browser download - log.info(f" [CARSI] Trying browser download for {doi}...") - result = client.download_via_browser(doi, resolved_url, output_path) + result = client.download_via_browser( + doi, + resolved_url, + output_path, + cancel_event=cancel_event, + ) if result: return result except Exception as e: - log.info(f" [CARSI] {e}") + if not _cancelled(cancel_event): + log.info(f" [CARSI] {e}") + finally: + if client is not None: + client.close() return None -def _save_pdf_response(resp: requests.Response, output_path: Path, doi: str, source: str) -> dict[str, Any] | None: +def _save_pdf_response( + resp: requests.Response, + output_path: Path, + doi: str, + source: str, + cancel_event: threading.Event | None = None, +) -> dict[str, Any] | None: """Save a PDF response to disk and validate it.""" try: + if _cancelled(cancel_event): + return None iterator = resp.iter_content(chunk_size=8192) first = next(iterator, b"") + if _cancelled(cancel_event): + return None if not _response_looks_pdf(resp, first): return None - output_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = output_path.with_suffix(output_path.suffix + ".part") - try: - with tmp_path.open("wb") as fh: - fh.write(first) - for chunk in iterator: - if chunk: - fh.write(chunk) - tmp_path.replace(output_path) - except Exception: - tmp_path.unlink(missing_ok=True) - raise - if is_pdf_file(output_path): + if not write_pdf_stream_atomic( + output_path, + first, + iterator, + cancel_event, + ): + return None + if not _cancelled(cancel_event) and is_pdf_file(output_path): return success(doi, output_path, source) except Exception: pass return None -def _try_browser_via_webvpn(url: str, config: dict[str, Any]) -> str | None: +def _try_browser_via_webvpn( + url: str, + config: dict[str, Any], + cancel_event: threading.Event | None = None, +) -> str | None: """Try fetching a URL through CloakBrowser, using WebVPN proxy.""" + if _cancelled(cancel_event): + return None base = _get_webvpn_base(config) proxied_url = convert_url(url, base, config) try: from ..browser_engine import is_available as browser_avail, get_html as browser_html if browser_avail(config): + if _cancelled(cancel_event): + return None result = browser_html(proxied_url, config) - if result: + if result and not _cancelled(cancel_event): return result except Exception as e: log.info(f" [browser] {e}") @@ -1169,38 +1534,44 @@ def _download_pdf_instsci( output_path: Path, config: dict[str, Any], doi: str, + cancel_event: threading.Event | None = None, ) -> dict[str, Any] | None: - if not is_plausible_pdf_url(url): + if _cancelled(cancel_event) or not is_plausible_pdf_url(url): return None + resp = None try: - _instsci_rate_limit() + if not _instsci_rate_limit(cancel_event): + return None resp = _fetch_via_webvpn(url, config, stream=True) - if resp.status_code >= 400: + if _cancelled(cancel_event) or resp.status_code >= 400: return None iterator = resp.iter_content(chunk_size=8192) first_chunk = next(iterator, b"") + if _cancelled(cancel_event): + return None if not _response_looks_pdf(resp, first_chunk): return None - output_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = output_path.with_suffix(output_path.suffix + ".part") - try: - with tmp_path.open("wb") as fh: - fh.write(first_chunk) - for chunk in iterator: - if chunk: - fh.write(chunk) - tmp_path.replace(output_path) - except Exception: - tmp_path.unlink(missing_ok=True) - raise + if not write_pdf_stream_atomic( + output_path, + first_chunk, + iterator, + cancel_event, + ): + return None - if is_pdf_file(output_path): + if not _cancelled(cancel_event) and is_pdf_file(output_path): result = success(doi, output_path, "WebVPN") result["doi"] = doi result["identifier"] = doi return result except Exception: pass + finally: + if resp is not None: + try: + resp.close() + except Exception: + pass return None diff --git a/src/scansci_pdf/sources/publishers.py b/src/scansci_pdf/sources/publishers.py index 53e34b2..03b6b3b 100644 --- a/src/scansci_pdf/sources/publishers.py +++ b/src/scansci_pdf/sources/publishers.py @@ -289,60 +289,115 @@ def try_publisher_direct(doi: str, output_path: Path, config: dict[str, Any]) -> # ============================================================ def try_mdpi_direct(doi: str, output_path: Path, config: dict[str, Any]) -> dict[str, Any] | None: - """Download MDPI open access papers directly.""" + """Download an MDPI paper only through its official PDF routes.""" if not doi.startswith("10.3390/"): return None + from ..network import USER_AGENT, fetch_json, polite_delay from ..pdf_utils import is_pdf_file, success, _response_looks_pdf - from ..network import USER_AGENT, polite_delay + from ..publisher_pdf_router import ( + _mdpi_landing_url_from_doi, + _mdpi_pdf_url_from_landing_url, + ) - article_id = doi.split("10.3390/")[-1] + urls: list[str] = [] - urls = [ - f"https://www.mdpi.com/{article_id}/pdf", - ] + def add_url(url: str | None) -> None: + candidate = str(url or "").strip() + if candidate and candidate not in urls: + urls.append(candidate) + try: + payload = fetch_json( + f"https://api.crossref.org/works/{requests.utils.quote(doi, safe='')}", + config, + headers={"Accept": "application/json"}, + ) + message = payload.get("message", {}) if isinstance(payload, dict) else {} + links = message.get("link", []) + if isinstance(links, dict): + links = [links] + if isinstance(links, list): + for link in links: + if not isinstance(link, dict): + continue + link_url = str(link.get("URL") or "") + if "mdpi.com/" in link_url.lower(): + add_url(link_url) + + landing_url = str( + message.get("resource", {}).get("primary", {}).get("URL") + or message.get("URL") + or "" + ) + add_url(_mdpi_pdf_url_from_landing_url(landing_url)) + except Exception as exc: + log.info(f" [MDPI] Crossref route lookup failed: {exc}") + + add_url(_mdpi_pdf_url_from_landing_url(_mdpi_landing_url_from_doi(doi))) + if not urls: + log.info(f" [MDPI] No official PDF route found for {doi}") + return None + + session = None try: session = requests.Session() session.trust_env = False session.headers.update({"User-Agent": USER_AGENT}) polite_delay(config) - # First try landing page to get real PDF URL - try: - resp = session.get(f"https://www.mdpi.com/{article_id}", timeout=10, allow_redirects=True) - if resp.status_code == 200: - pdf_match = re.search(r'citation_pdf_url["\s]+content="([^"]+)"', resp.text[:5000], re.I) - if not pdf_match: - pdf_match = re.search(r'href="(/[^"]*?/pdf[^"]*)"', resp.text[:10000], re.I) - if pdf_match: - pdf_url = pdf_match.group(1) - if pdf_url.startswith("/"): - pdf_url = f"https://www.mdpi.com{pdf_url}" - urls.insert(0, pdf_url) - except Exception: - pass - for pdf_url in urls: + response = None try: - resp2 = session.get(pdf_url, timeout=15, stream=True, - headers={"Accept": "application/pdf,*/*"}) - if resp2.status_code >= 400: + response = session.get( + pdf_url, + timeout=15, + stream=True, + headers={"Accept": "application/pdf,*/*"}, + ) + if response.status_code >= 400: + log.info(f" [MDPI] HTTP {response.status_code}: {pdf_url}") continue - iterator = resp2.iter_content(chunk_size=8192) + iterator = response.iter_content(chunk_size=8192) first = next(iterator, b"") - if not _response_looks_pdf(resp2, first): + if not _response_looks_pdf(response, first): continue - if not _write_pdf_atomic(output_path, first, iterator): continue if is_pdf_file(output_path): return success(doi, output_path, "MDPIDirect") except Exception: continue - except Exception: - pass + finally: + if response is not None: + try: + response.close() + except Exception: + pass + finally: + if session is not None: + try: + session.close() + except Exception: + pass + + if config.get("browser_enabled", True): + try: + from ..browser_engine import download_pdf_via_browser, is_available + + if is_available(config): + for pdf_url in urls: + log.info(f" [MDPI] browser download: {pdf_url}") + if download_pdf_via_browser( + pdf_url, + output_path, + config, + timeout=90.0, + ) and is_pdf_file(output_path): + return success(doi, output_path, "MDPI(Browser)") + except Exception as exc: + log.info(f" [MDPI] browser fallback failed: {exc}") return None