fix(browser): preserve user downloads when attached

sentient/browser/service.py

@@ -69,6 +69,7 @@
 DOWNLOAD_CANCEL_TIMEOUT_S = 5.0
 DOWNLOAD_TASK_RETENTION_S = 300.0
 DOWNLOAD_TASK_CLEANUP_INTERVAL_S = 30.0
+ATTACHED_MATCH_TIMEOUT_S = 2.0
 # models often pass the whole snapshot line ("[e4] button \"Place order\"") instead of just "e4"
 _REF_RE = re.compile(r"\b(e\d+)\b", re.IGNORECASE)
 
@@ -238,6 +239,9 @@ def __init__(self, app):
         self._download_completed_at: dict[asyncio.Task[str], float] = {}
         self._download_cleanup_task: asyncio.Task | None = None
         self._download_lock = asyncio.Lock()
+        self._browser_cdp: Any = None
+        self._attached_downloads: list[dict[str, Any]] = []
+        self._attached_match_timeout_s = ATTACHED_MATCH_TIMEOUT_S
 
     # ------------------------------------------------------------------ lifecycle
     async def start(self) -> None:
@@ -602,6 +606,12 @@ async def _attach(self, prof: Any) -> Any:
         self._browser = browser
         self._attached = True
         self._engine = None
+        with contextlib.suppress(Exception):
+            cdp = await browser.new_browser_cdp_session()
+            self._browser_cdp = cdp
+            cdp.on("Browser.downloadWillBegin", self._on_cdp_download_will_begin)
+            cdp.on("Browser.downloadProgress", self._on_cdp_download_progress)
+            await cdp.send("Browser.setDownloadBehavior", {"behavior": "default", "eventsEnabled": True})
         return context
 
     async def _ensure(self, ctx: Any = None) -> Any:
@@ -657,6 +667,11 @@ async def _close_context(self) -> None:
             self._focused = None
             self._download_pages.clear()
             self._internal_pages.clear()
+            if self._browser_cdp is not None:
+                with contextlib.suppress(Exception):
+                    await self._browser_cdp.detach()
+                self._browser_cdp = None
+            self._clear_attached_downloads()
 
     async def _flush_storage(self, ctx: Any) -> None:
         """Let a launched browser write site storage to disk before it closes: close its tabs with their unload
@@ -750,6 +765,8 @@ def _on_context_closed(self, context: Any) -> None:
         self._snap = None
         self._focused = None
         self._download_pages.clear()
+        self._browser_cdp = None
+        self._clear_attached_downloads()
         if not self._closing:  # the user closed the visible window (or their attached browser)
             self._spawn(self._publish_status())
 
@@ -865,9 +882,132 @@ async def _cleanup_download_tasks(self) -> None:
         finally:
             self._download_cleanup_task = None
 
+    def _clear_attached_downloads(self, exc: Exception | None = None) -> None:
+        """Safely fail and clear all tracked attached-browser CDP downloads during context shutdown."""
+        entries = list(self._attached_downloads)
+        self._attached_downloads.clear()
+        err = exc or BrowserError("The browser closed before the download completed.")
+        for entry in entries:
+            fut = entry.get("future")
+            if fut is not None and not fut.done():
+                fut.set_exception(err)
+                with contextlib.suppress(Exception):
+                    fut.exception()
+
+    def _prune_attached_downloads(self) -> None:
+        """Evict completed download records while keeping pending downloads reachable."""
+        if len(self._attached_downloads) <= 100:
+            return
+        retained: list[dict[str, Any]] = []
+        for entry in self._attached_downloads:
+            fut = entry.get("future")
+            if (fut is not None and not fut.done()) or not entry.get("consumed"):
+                retained.append(entry)
+        if len(retained) > 100:
+            pending = [e for e in retained if e.get("future") is not None and not e["future"].done()]
+            finished = [e for e in retained if e.get("future") is None or e["future"].done()]
+            max_finished = max(0, 100 - len(pending))
+            retained = pending + finished[-max_finished:]
+        if len(retained) > 100:
+            to_evict = retained[:-100]
+            retained = retained[-100:]
+            for entry in to_evict:
+                fut = entry.get("future")
+                if fut is not None and not fut.done():
+                    fut.set_exception(BrowserError("Download record evicted."))
+                    with contextlib.suppress(Exception):
+                        fut.exception()
+        self._attached_downloads = retained
+
+    def _on_cdp_download_will_begin(self, payload: dict[str, Any]) -> None:
+        """Record a CDP Browser.downloadWillBegin event to attribute downloads to frames."""
+        guid = payload.get("guid")
+        if not guid:
+            return
+        entry = {
+            "guid": guid,
+            "frameId": payload.get("frameId", ""),
+            "url": payload.get("url", ""),
+            "filename": self._download_filename(payload.get("suggestedFilename") or "download"),
+            "future": asyncio.get_running_loop().create_future(),
+            "consumed": False,
+        }
+        self._attached_downloads.append(entry)
+        self._prune_attached_downloads()
+
+    def _on_cdp_download_progress(self, payload: dict[str, Any]) -> None:
+        """Handle CDP Browser.downloadProgress events to complete or fail awaiting download tasks."""
+        guid = payload.get("guid")
+        state = payload.get("state")
+        file_path = payload.get("filePath")
+        if not guid:
+            return
+        for entry in self._attached_downloads:
+            if entry["guid"] == guid:
+                fut = entry.get("future")
+                if fut is not None and not fut.done():
+                    if state == "completed" and file_path:
+                        fut.set_result(file_path)
+                    elif state == "canceled":
+                        fname = entry.get("filename") or "download"
+                        fut.set_exception(BrowserError(f"Download '{fname}' was canceled by the browser."))
+                        with contextlib.suppress(Exception):
+                            fut.exception()
+                break
+
+    @staticmethod
+    def _match_cdp_entry(
+        entries: list[dict[str, Any]],
+        page_frames: set[str],
+        dl_url: str,
+        dl_name: str,
+    ) -> dict[str, Any] | None:
+        """Find the best matching unconsumed CDP download entry for a given frame tree and download metadata."""
+        best_entry: dict[str, Any] | None = None
+        best_score = 0
+        for entry in entries:
+            if entry.get("frameId") not in page_frames or entry.get("consumed"):
+                continue
+            e_url = entry.get("url", "")
+            e_name = entry.get("filename", "")
+
+            url_matches = bool(dl_url and e_url and dl_url == e_url)
+            name_matches = bool(dl_name and e_name and dl_name not in {"", "download"} and dl_name == e_name)
+            url_conflicts = bool(dl_url and e_url and dl_url != e_url)
+            name_conflicts = bool(
+                dl_name
+                and e_name
+                and dl_name not in {"", "download"}
+                and e_name not in {"", "download"}
+                and dl_name != e_name
+            )
+
+            if url_conflicts:
+                continue
+            if name_conflicts and not url_matches:
+                continue
+
+            score = 0
+            if url_matches and name_matches:
+                score = 3
+            elif url_matches or name_matches:
+                score = 2
+            elif not name_conflicts:
+                score = 1
+
+            if score > best_score:
+                best_score = score
+                best_entry = entry
+                if score == 3:
+                    break
+        return best_entry
+
     @staticmethod
     def _download_filename(download: Any) -> str:
-        raw_name = str(getattr(download, "suggested_filename", None) or "download")
+        if isinstance(download, str):
+            raw_name = download
+        else:
+            raw_name = str(getattr(download, "suggested_filename", None) or "download")
         name = raw_name.replace("\\", "/").rsplit("/", 1)[-1]
         return re.sub(r"[^\w.\- ()]+", "_", name).strip(" .")[:160] or "download"
 
@@ -918,12 +1058,65 @@ def save_failure(exc: BaseException) -> str:
             reason = str(exc).strip().splitlines()[0][:300] or type(exc).__name__
             return f"Couldn't save downloaded file '{name}': {reason}"
 
+        def abort_closed_browser() -> None:
+            raise BrowserError("The browser closed before the download completed.")
+
         timeout = asyncio.timeout(DOWNLOAD_MAX_DURATION_S)
         target: Path | None = None
         try:
             async with timeout:
                 target = await reserve_target()
-                await download.save_as(str(target))
+                if self._attached:
+                    matched = None
+                    page = getattr(download, "page", None)
+                    if page is not None and self._context is not None:
+                        page_frames: set[str] = set()
+                        try:
+                            page_cdp = await self._context.new_cdp_session(page)
+                            try:
+                                tree = await page_cdp.send("Page.getFrameTree")
+
+                                def _collect_frame_ids(node: dict[str, Any]) -> set[str]:
+                                    fids = {node.get("frame", {}).get("id", "")}
+                                    for child in node.get("childFrames", []):
+                                        fids.update(_collect_frame_ids(child))
+                                    return fids
+
+                                page_frames = _collect_frame_ids(tree.get("frameTree", {}))
+                            finally:
+                                with contextlib.suppress(Exception):
+                                    await page_cdp.detach()
+                        except Exception:
+                            pass
+
+                        if page_frames:
+                            match_timeout = getattr(self, "_attached_match_timeout_s", ATTACHED_MATCH_TIMEOUT_S)
+                            match_deadline = time.monotonic() + match_timeout
+                            dl_url = str(getattr(download, "url", "") or "")
+                            while True:
+                                matched = self._match_cdp_entry(
+                                    self._attached_downloads, page_frames, dl_url, name
+                                )
+                                if matched is not None or time.monotonic() >= match_deadline or self._context is None:
+                                    break
+                                await asyncio.sleep(0.05)
+
+                    if matched is not None:
+                        matched["consumed"] = True
+                        file_path = await matched["future"]
+
+                        def _finalize_file() -> None:
+                            assert target is not None
+                            target.unlink(missing_ok=True)
+                            shutil.move(file_path, str(target))
+
+                        await asyncio.to_thread(_finalize_file)
+                    elif self._context is None:
+                        abort_closed_browser()
+                    else:
+                        await download.save_as(str(target))
+                else:
+                    await download.save_as(str(target))
         except asyncio.CancelledError as exc:
             try:
                 cancel_error = await request_download_cancel()