1"""Fetch transport: DIRECT MODE first (httpx, conditional requests, per-domain rate limits, robots.txt, backoff), optional2ESCALATED MODE (headless browser → Scrapfly → Firecrawl) that is never required and never automatic unless the connector asks.34 direct → retry → [browser] → [scrapfly] → [firecrawl] → BlockedError (→ review queue by the caller)5"""6from __future__ import annotations78import asyncio9import hashlib10import logging11import time12from dataclasses import dataclass, field13from datetime import UTC, datetime14from urllib.parse import urlparse, urlunparse15from urllib.robotparser import RobotFileParser1617import httpx1819from aiatlas.config import settings2021log = logging.getLogger(__name__)2223TRANSIENT_STATUS = {408, 425, 429, 500, 502, 503, 504}24BLOCK_STATUS = {401, 403, 451, 999}252627class FetchError(Exception):28 def __init__(self, message: str, *, status: int | None = None, url: str = ""):29 super().__init__(message)30 self.status = status31 self.url = url323334class BlockedError(FetchError):35 """Access denied by the origin after every allowed transport (never bypass: park in the review queue)."""363738class NotModified(Exception):39 """HTTP 304 — content unchanged since our stored validators."""404142@dataclass43class FetchResult:44 url: str45 final_url: str46 status: int47 headers: dict[str, str]48 content: bytes49 content_type: str50 fetched_at: datetime51 duration_ms: int52 transport: str = "direct"53 from_cache: bool = False54 sha256: str = field(init=False)5556 def __post_init__(self) -> None:57 self.sha256 = hashlib.sha256(self.content).hexdigest()5859 @property60 def text(self) -> str:61 enc = "utf-8"62 ct = self.content_type.lower()63 if "charset=" in ct:64 enc = ct.split("charset=", 1)[1].split(";")[0].strip().strip('"') or "utf-8"65 try:66 return self.content.decode(enc, errors="replace")67 except LookupError:68 return self.content.decode("utf-8", errors="replace")6970 @property71 def is_html(self) -> bool:72 return "html" in self.content_type or self.content[:256].lstrip().lower().startswith((b"<!doctype html", b"<html"))7374 @property75 def is_json(self) -> bool:76 return "json" in self.content_type or self.content[:1] in (b"{", b"[")7778 @property79 def is_xml(self) -> bool:80 return "xml" in self.content_type or self.content[:5] == b"<?xml"8182 @property83 def is_pdf(self) -> bool:84 return "pdf" in self.content_type or self.content[:4] == b"%PDF"8586 def json(self): # type: ignore[no-untyped-def]87 import orjson8889 return orjson.loads(self.content)909192def canonicalize_url(url: str) -> str:93 """Stable URL key: lowercase scheme/host, no fragment, no tracking params, no trailing slash duplication."""94 p = urlparse(url.strip())95 query = "&".join(sorted(q for q in p.query.split("&") if q and not q.lower().startswith(96 ("utm_", "ref=", "ref_", "fbclid", "gclid", "mc_cid", "mc_eid", "_hs", "igshid", "source="))))97 path = p.path or "/"98 if len(path) > 1 and path.endswith("/") and not path.endswith("//"):99 path = path.rstrip("/") or "/"100 return urlunparse((p.scheme.lower() or "https", p.netloc.lower(), path, "", query, ""))101102103def domain_of(url: str) -> str:104 host = urlparse(url).netloc.lower()105 return host[4:] if host.startswith("www.") else host106107108class _RateLimiter:109 """Per-domain minimum spacing between requests (token bucket, in-process). Cluster-wide fairness comes from Redis locks110 at the connector level (one connector run at a time)."""111112 def __init__(self) -> None:113 self._next: dict[str, float] = {}114 self._locks: dict[str, asyncio.Lock] = {}115116 async def wait(self, domain: str, per_min: int) -> None:117 lock = self._locks.setdefault(domain, asyncio.Lock())118 async with lock:119 spacing = 60.0 / max(1, per_min)120 now = time.monotonic()121 ready = self._next.get(domain, 0.0)122 if ready > now:123 await asyncio.sleep(ready - now)124 now = time.monotonic()125 self._next[domain] = now + spacing126127128class _Robots:129 def __init__(self) -> None:130 self._cache: dict[str, tuple[float, RobotFileParser | None]] = {}131132 async def allowed(self, client: httpx.AsyncClient, url: str) -> bool:133 if not settings.respect_robots:134 return True135 p = urlparse(url)136 key = f"{p.scheme}://{p.netloc}"137 cached = self._cache.get(key)138 if cached is None or cached[0] < time.monotonic():139 rp: RobotFileParser | None = RobotFileParser()140 try:141 r = await client.get(f"{key}/robots.txt", timeout=15, follow_redirects=True)142 if r.status_code == 200 and len(r.content) < 512 * 1024:143 rp.parse(r.text.splitlines()) # type: ignore[union-attr]144 else:145 rp = None # no robots or error → allowed146 except Exception: # noqa: BLE001147 rp = None148 self._cache[key] = (time.monotonic() + 6 * 3600, rp)149 cached = self._cache[key]150 rp = cached[1]151 if rp is None:152 return True153 try:154 return rp.can_fetch(settings.user_agent.split("/")[0], url) or rp.can_fetch("*", url)155 except Exception: # noqa: BLE001156 return True157158159class Fetcher:160 """One Fetcher per connector run. Direct mode by default; `escalate=True` enables the optional chain."""161162 def __init__(self, *, rate_per_min: int | None = None, robots: bool = True, timeout_s: float | None = None,163 headers: dict[str, str] | None = None, http2: bool = True):164 self.rate_per_min = rate_per_min or settings.default_rate_per_min165 self.robots = robots166 self.timeout_s = timeout_s or settings.http_timeout_s167 self.headers = {"User-Agent": settings.user_agent, "Accept": "text/html,application/xhtml+xml,application/xml,application/json,text/*;q=0.9,*/*;q=0.8",168 "Accept-Language": "en-US,en;q=0.8", **(headers or {})}169 self._client: httpx.AsyncClient | None = None170 self._http2 = http2171172 async def __aenter__(self) -> Fetcher:173 self._client = httpx.AsyncClient(headers=self.headers, timeout=httpx.Timeout(self.timeout_s, connect=20), follow_redirects=True,174 http2=self._http2, limits=httpx.Limits(max_connections=16, max_keepalive_connections=8))175 return self176177 async def __aexit__(self, *exc: object) -> None:178 if self._client:179 await self._client.aclose()180 self._client = None181182 @property183 def client(self) -> httpx.AsyncClient:184 if self._client is None:185 raise RuntimeError("use `async with Fetcher() as f`")186 return self._client187188 async def get(self, url: str, *, etag: str | None = None, last_modified: str | None = None, min_bytes: int = 1,189 escalate: bool = False, retries: int = 2, accept: str | None = None, rate_per_min: int | None = None) -> FetchResult:190 """Direct fetch with conditional headers. Raises NotModified (304), FetchError, BlockedError."""191 client = self.client192 domain = domain_of(url)193 if self.robots and not await _robots.allowed(client, url):194 raise BlockedError(f"robots.txt disallows {url}", status=None, url=url)195 headers: dict[str, str] = {}196 if etag:197 headers["If-None-Match"] = etag198 if last_modified:199 headers["If-Modified-Since"] = last_modified200 if accept:201 headers["Accept"] = accept202 last_exc: Exception | None = None203 for attempt in range(retries + 1):204 await _limiter.wait(domain, rate_per_min or self.rate_per_min)205 t0 = time.perf_counter()206 try:207 async with client.stream("GET", url, headers=headers) as r:208 if r.status_code == 304:209 raise NotModified()210 if r.status_code in TRANSIENT_STATUS and attempt < retries:211 retry_after = r.headers.get("retry-after")212 delay = min(60.0, float(retry_after)) if retry_after and retry_after.isdigit() else 2.0 * (2 ** attempt)213 log.info("transient status, backing off", extra={"url": url, "status": r.status_code, "delay": delay})214 await asyncio.sleep(delay)215 continue216 if r.status_code in BLOCK_STATUS:217 if escalate:218 return await self._escalate(url, reason=f"http {r.status_code}")219 raise BlockedError(f"http {r.status_code} for {url}", status=r.status_code, url=url)220 if r.status_code >= 400:221 raise FetchError(f"http {r.status_code} for {url}", status=r.status_code, url=url)222 chunks: list[bytes] = []223 size = 0224 async for chunk in r.aiter_bytes():225 size += len(chunk)226 if size > settings.max_body_bytes:227 raise FetchError(f"body exceeds {settings.max_body_bytes} bytes", status=r.status_code, url=url)228 chunks.append(chunk)229 content = b"".join(chunks)230 if len(content) < min_bytes:231 raise FetchError(f"suspiciously short body ({len(content)} bytes) for {url}", status=r.status_code, url=url)232 res = FetchResult(url=url, final_url=str(r.url), status=r.status_code, headers={k.lower(): v for k, v in r.headers.items()},233 content=content, content_type=r.headers.get("content-type", "").lower(), fetched_at=datetime.now(UTC),234 duration_ms=int((time.perf_counter() - t0) * 1000), transport="direct")235 if escalate and res.is_html and _looks_like_challenge(content):236 return await self._escalate(url, reason="anti-bot challenge page")237 return res238 except NotModified:239 raise240 except (BlockedError, FetchError):241 raise242 except (httpx.TimeoutException, httpx.TransportError) as exc:243 last_exc = exc244 if attempt < retries:245 await asyncio.sleep(1.5 * (2 ** attempt))246 continue247 raise FetchError(f"{exc.__class__.__name__}: {exc}", url=url) from exc248 raise FetchError(f"fetch failed: {last_exc}", url=url)249250 # ---------------------------------------------------------------------------------------------- escalation (optional)251 async def _escalate(self, url: str, *, reason: str) -> FetchResult:252 log.warning("escalating fetch", extra={"url": url, "reason": reason})253 if settings.browser_enabled:254 try:255 return await self._browser(url)256 except Exception as exc: # noqa: BLE001257 log.info("browser transport failed", extra={"url": url, "error": str(exc)})258 if settings.scrapfly_api_key:259 try:260 return await self._scrapfly(url)261 except Exception as exc: # noqa: BLE001262 log.info("scrapfly transport failed", extra={"url": url, "error": str(exc)})263 if settings.firecrawl_api_key:264 try:265 return await self._firecrawl(url)266 except Exception as exc: # noqa: BLE001267 log.info("firecrawl transport failed", extra={"url": url, "error": str(exc)})268 raise BlockedError(f"blocked after escalation ({reason}): {url}", url=url)269270 async def _browser(self, url: str) -> FetchResult:271 from playwright.async_api import async_playwright # optional dependency272273 t0 = time.perf_counter()274 async with async_playwright() as p:275 browser = await p.chromium.launch(headless=True)276 try:277 page = await browser.new_page(user_agent=settings.user_agent)278 resp = await page.goto(url, wait_until="networkidle", timeout=int(self.timeout_s * 1000))279 html = await page.content()280 status = resp.status if resp else 200281 finally:282 await browser.close()283 return FetchResult(url=url, final_url=url, status=status, headers={}, content=html.encode(), content_type="text/html; charset=utf-8",284 fetched_at=datetime.now(UTC), duration_ms=int((time.perf_counter() - t0) * 1000), transport="browser")285286 async def _scrapfly(self, url: str) -> FetchResult:287 t0 = time.perf_counter()288 r = await self.client.get("https://api.scrapfly.io/scrape", params={"key": settings.scrapfly_api_key, "url": url, "asp": "true",289 "render_js": "true", "country": "us"}, timeout=120)290 r.raise_for_status()291 data = r.json()["result"]292 content = (data.get("content") or "").encode()293 return FetchResult(url=url, final_url=data.get("url") or url, status=int(data.get("status_code") or 200), headers={},294 content=content, content_type=(data.get("content_type") or "text/html").lower(), fetched_at=datetime.now(UTC),295 duration_ms=int((time.perf_counter() - t0) * 1000), transport="scrapfly")296297 async def _firecrawl(self, url: str) -> FetchResult:298 t0 = time.perf_counter()299 r = await self.client.post("https://api.firecrawl.dev/v1/scrape", json={"url": url, "formats": ["rawHtml"]},300 headers={"Authorization": f"Bearer {settings.firecrawl_api_key}"}, timeout=120)301 r.raise_for_status()302 data = r.json().get("data", {})303 content = (data.get("rawHtml") or data.get("html") or "").encode()304 return FetchResult(url=url, final_url=(data.get("metadata") or {}).get("sourceURL") or url, status=200, headers={}, content=content,305 content_type="text/html", fetched_at=datetime.now(UTC), duration_ms=int((time.perf_counter() - t0) * 1000),306 transport="firecrawl")307308309def _looks_like_challenge(content: bytes) -> bool:310 head = content[:20000].lower()311 if len(content) > 60_000:312 return False313 markers = (b"just a moment", b"cf-chl-", b"challenge-platform", b"attention required", b"verify you are human", b"access denied",314 b"captcha", b"perimeterx", b"_px", b"datadome", b"enable javascript and cookies to continue")315 return sum(m in head for m in markers) >= 2316317318def file_result(path: str, *, url: str, content_type: str = "application/octet-stream") -> FetchResult:319 """Wrap a local file as a FetchResult (fixtures, seed snapshots, `aia run x --file`)."""320 with open(path, "rb") as fh:321 content = fh.read()322 return FetchResult(url=url, final_url=url, status=200, headers={}, content=content, content_type=content_type,323 fetched_at=datetime.now(UTC), duration_ms=0, transport="file")324325326_limiter = _RateLimiter()327_robots = _Robots()328329__all__ = ["Fetcher", "FetchResult", "FetchError", "BlockedError", "NotModified", "canonicalize_url", "domain_of", "file_result"]330