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,58 @@ async def _cleanup_download_tasks(self) -> None:
         finally:
             self._download_cleanup_task = None
 
+    def _clear_attached_downloads(self, exc: Exception | None = None) -> None:
+        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 _on_cdp_download_will_begin(self, payload: dict[str, Any]) -> None:
+        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)
+        if len(self._attached_downloads) > 100:
+            self._attached_downloads = [e for e in self._attached_downloads if not e.get("consumed")][-50:]
+
+    def _on_cdp_download_progress(self, payload: dict[str, Any]) -> None:
+        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 _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 +984,60 @@ 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
+                            while True:
+                                for entry in self._attached_downloads:
+                                    if entry.get("frameId") in page_frames and not entry.get("consumed"):
+                                        matched = entry
+                                        break
+                                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"]
+                        target.unlink(missing_ok=True)
+                        shutil.move(file_path, str(target))
+                    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()