Screenshots: 2-phase (HTTPS-proxy resolve + plain media-proxy chunked download + local ffmpeg)

pipeline/frames.py

@@ -1,16 +1,18 @@
-"""Stage 5 + 6: real screenshots via short yt-dlp clip downloads + ffmpeg, at weighted
-timestamps.
-
-For each timestamp we download a ~1.5 s clip with yt-dlp's ``download_ranges`` (a couple
-hundred KB) and let ``ffmpeg`` extract frames from that *local* file. We grab 3 candidates
-around the target and keep the sharpest (Laplacian variance) to avoid a blurry/transition
-frame. Downloading (rather than ffmpeg-seeking the remote stream) is required because the
-stream lives on ``googlevideo.com`` and must go through the residential proxy — and when
-``YT_PROXY`` is an HTTPS/TLS-wrapped proxy (needed to hide the target host from egress DPI)
-ffmpeg can't talk to it, but yt-dlp can.
-
-The capture timestamp itself comes from the weighted indicator: a blend of the LLM's
-suggested timestamp and the transcript segment timing of the text the step quotes.
+"""Stage 5 + 6: real screenshots via a one-time low-res video download + ffmpeg, at
+weighted timestamps.
+
+Two proxies are used because of how the Space's egress is filtered:
+
+* Resolution hits ``youtube.com`` and must go through ``YT_PROXY`` — an HTTPS/TLS-wrapped
+  proxy that hides the target host from the egress DPI that otherwise resets YouTube.
+* The media itself lives on ``googlevideo.com`` (not DPI-filtered) and is many MB, which
+  the TLS proxy can't sustain — so it's downloaded through ``YT_MEDIA_PROXY``, a plain
+  HTTP proxy, in chunked Range requests (dodging YouTube's single-stream throttling).
+
+We download one capped-resolution copy, then ``ffmpeg`` reads the *local* file to grab 3
+candidates around each timestamp and keeps the sharpest (Laplacian variance). The capture
+timestamp comes from the weighted indicator: a blend of the LLM's suggested timestamp and
+the transcript segment timing of the text the step quotes.
 """
 from __future__ import annotations
 
@@ -18,8 +20,8 @@ import difflib
 import os
 import subprocess
 
+import requests
 from yt_dlp import YoutubeDL
-from yt_dlp.utils import download_range_func
 
 
 # --------------------------------------------------------------------------- weighting
@@ -84,15 +86,79 @@ def _sharpness(path: str) -> float:
     return float(lap.var())
 
 
-def _best_frame_from_clip(clip: str, out_dir: str, idx: int, target: float) -> str | None:
-    """Extract candidates near ``target`` seconds into a *local* ``clip`` and keep the
+def _resolve_media_url(video_id: str, max_height: int, cookiefile: str | None,
+                       proxy: str | None) -> str | None:
+    """Resolve a single-file (progressive/DASH) video URL <= ``max_height`` via the
+    HTTPS ``proxy``. Returns the ``googlevideo.com`` media URL, or None."""
+    opts = {
+        "quiet": True, "no_warnings": True, "skip_download": True,
+        "format": (f"bv*[height<={max_height}][ext=mp4]/bv*[height<={max_height}]/"
+                   f"best[height<={max_height}]/best"),
+    }
+    if cookiefile:
+        opts["cookiefile"] = cookiefile
+    if proxy:
+        opts["proxy"] = proxy
+        # yt-dlp validates against certifi only and ignores SSL_CERT_FILE; switch it to
+        # the default cert path (which honors our combined bundle) so the self-signed
+        # HTTPS-proxy cert validates. See app._install_proxy_ca.
+        if os.environ.get("SSL_CERT_FILE"):
+            opts["compat_opts"] = ["no-certifi"]
+    with YoutubeDL(opts) as ydl:
+        info = ydl.extract_info(f"https://www.youtube.com/watch?v={video_id}", download=False)
+    if info.get("url"):
+        return info["url"]
+    vids = [f for f in info.get("formats", [])
+            if f.get("vcodec") not in (None, "none") and f.get("url")
+            and (f.get("height") or 0) <= max_height
+            and str(f.get("protocol", "")).startswith("http")]
+    if vids:
+        return sorted(vids, key=lambda f: f.get("height") or 0)[-1]["url"]
+    return None
+
+
+def _download_media(url: str, dest: str, media_proxy: str | None,
+                    chunk: int = 1 << 20, timeout: int = 30) -> None:
+    """Download ``url`` to ``dest`` through the plain ``media_proxy`` using chunked Range
+    requests — small ranges dodge YouTube's single-stream throttling and stay within the
+    plain proxy's throughput (the TLS proxy can't sustain large transfers)."""
+    proxies = {"http": media_proxy, "https": media_proxy} if media_proxy else None
+    probe = requests.get(url, proxies=proxies, headers={"Range": "bytes=0-0"}, timeout=timeout)
+    probe.raise_for_status()
+    cr = probe.headers.get("Content-Range", "")
+    total = int(cr.split("/")[-1]) if "/" in cr and cr.split("/")[-1].isdigit() else 0
+    with open(dest, "wb") as fh:
+        if not total:  # server ignored Range — fall back to a single stream
+            with requests.get(url, proxies=proxies, stream=True, timeout=timeout) as r:
+                r.raise_for_status()
+                for c in r.iter_content(chunk):
+                    if c:
+                        fh.write(c)
+            return
+        start = 0
+        while start < total:
+            end = min(start + chunk - 1, total - 1)
+            for attempt in range(3):
+                try:
+                    rr = requests.get(url, proxies=proxies, timeout=timeout,
+                                      headers={"Range": f"bytes={start}-{end}"})
+                    rr.raise_for_status()
+                    fh.write(rr.content)
+                    break
+                except requests.RequestException:
+                    if attempt == 2:
+                        raise
+            start = end + 1
+
+
+def _grab_local(media_path: str, t: float, out_dir: str, idx: int) -> str | None:
+    """Grab 3 candidates around ``t`` from the *local* ``media_path`` and keep the
     sharpest. ffmpeg reads a local file here, so no proxy/network is involved."""
     cands = []
-    for k, dt in enumerate((-0.4, 0.0, 0.4)):
+    for k, dt in enumerate((-0.5, 0.0, 0.5)):
         p = os.path.join(out_dir, f"shot_{idx}_{k}.jpg")
-        ss = max(0.0, target + dt)
-        cmd = ["ffmpeg", "-y", "-loglevel", "error", "-ss", f"{ss:.2f}",
-               "-i", clip, "-frames:v", "1", "-q:v", "2", p]
+        cmd = ["ffmpeg", "-y", "-loglevel", "error", "-ss", f"{max(0.0, t + dt):.2f}",
+               "-i", media_path, "-frames:v", "1", "-q:v", "2", p]
         try:
             subprocess.run(cmd, capture_output=True, text=True, timeout=60)
         except subprocess.TimeoutExpired:
@@ -113,51 +179,32 @@ def _best_frame_from_clip(clip: str, out_dir: str, idx: int, target: float) -> s
 
 def capture_shots(times: dict[int, float], video_id: str, out_dir: str,
                   cookiefile: str | None = None, proxy: str | None = None,
-                  progress=None, *, half: float = 0.75,
-                  max_height: int = 480) -> dict[int, dict]:
-    """Download a short clip around each ``{step_index: time}`` and keep the sharpest
-    frame; return ``{idx: {"time", "path"}}``. One yt-dlp call fetches every section."""
+                  media_proxy: str | None = None, progress=None, *,
+                  max_height: int = 360) -> dict[int, dict]:
+    """Download one capped-resolution copy of the video (resolve via ``proxy``, fetch
+    media via ``media_proxy``) and keep the sharpest frame per ``{step_index: time}``;
+    return ``{idx: {"time", "path"}}``."""
     os.makedirs(out_dir, exist_ok=True)
     items = list(times.items())
     if not items:
         return {}
-    clips_dir = os.path.join(out_dir, "_clips")
-    os.makedirs(clips_dir, exist_ok=True)
-    ranges = [(max(0.0, t - half), t + half) for _, t in items]
-    opts = {
-        "quiet": True, "no_warnings": True,
-        "format": (f"bv*[height<={max_height}][vcodec!=none]/"
-                   f"best[height<={max_height}]/best[vcodec!=none]/best"),
-        "paths": {"home": clips_dir},
-        "outtmpl": {"default": "sec_%(section_start)0.2f.%(ext)s"},
-        "download_ranges": download_range_func(None, ranges),
-        "force_keyframes_at_cuts": True,
-    }
-    if cookiefile:
-        opts["cookiefile"] = cookiefile
-    if proxy:
-        opts["proxy"] = proxy
-        # yt-dlp trusts only certifi's own bundle by default and ignores SSL_CERT_FILE.
-        # When an HTTPS-proxy CA has been installed into a combined bundle (SSL_CERT_FILE,
-        # set by app._install_proxy_ca), switch yt-dlp to the default cert path so it
-        # honors that file — this is what lets it validate the self-signed proxy cert.
-        if os.environ.get("SSL_CERT_FILE"):
-            opts["compat_opts"] = ["no-certifi"]
-    with YoutubeDL(opts) as ydl:
-        ydl.download([f"https://www.youtube.com/watch?v={video_id}"])
+    url = _resolve_media_url(video_id, max_height, cookiefile, proxy)
+    if not url:
+        raise RuntimeError("Could not resolve a downloadable video stream URL.")
+    media = os.path.join(out_dir, "_video.bin")
+    if progress:
+        progress(0.0, desc="Downloading video for screenshots")
+    _download_media(url, media, media_proxy or proxy)
 
-    import glob
     out: dict[int, dict] = {}
     for n, (idx, t) in enumerate(items):
         if progress:
             progress((n + 1) / max(1, len(items)), desc=f"Screenshot {n + 1}/{len(items)}")
-        start = max(0.0, t - half)
-        matches = [m for m in glob.glob(os.path.join(clips_dir, f"sec_{start:.2f}.*"))
-                   if not m.endswith(".part")]
-        if not matches:
-            continue
-        # t sits ``t - start`` seconds into its clip (== half, except when clamped at 0)
-        path = _best_frame_from_clip(matches[0], out_dir, idx, t - start)
+        path = _grab_local(media, t, out_dir, idx)
         if path:
             out[idx] = {"time": t, "path": path}
+    try:
+        os.remove(media)
+    except OSError:
+        pass
     return out