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._attached_downloads.clear()
 
     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._attached_downloads.clear()
         if not self._closing:  # the user closed the visible window (or their attached browser)
             self._spawn(self._publish_status())
 
@@ -865,9 +882,42 @@ async def _cleanup_download_tasks(self) -> None:
         finally:
             self._download_cleanup_task = None
 
+    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:
+                if state == "completed" and file_path and not entry["future"].done():
+                    entry["future"].set_result(file_path)
+                elif state == "canceled" and not entry["future"].done():
+                    entry["future"].cancel()
+                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"
 
@@ -923,7 +973,50 @@ def save_failure(exc: BaseException) -> str:
         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:
+                                    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))
+                    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()

tests/browser/test_profiles.py

@@ -6,6 +6,7 @@
 import http.server
 import json
 import threading
+import time
 from pathlib import Path
 
 import pytest
@@ -645,6 +646,231 @@ async def new_page(self):
     assert nested_popup in svc._download_pages
 
 
+async def test_attached_browser_resets_download_behavior_to_default(monkeypatch):
+    svc = BrowserService(make_app())
+    cdp_calls = []
+
+    class MockCDP:
+        def on(self, event, handler):
+            pass
+
+        async def send(self, command, params):
+            cdp_calls.append((command, params))
+
+    class MockBrowser:
+        contexts = [object()]
+
+        def on(self, event, handler):
+            pass
+
+        async def new_browser_cdp_session(self):
+            return MockCDP()
+
+    class MockChromium:
+        async def connect_over_cdp(self, ws, timeout=None):
+            return MockBrowser()
+
+    svc._pw = type("MockPW", (), {"chromium": MockChromium()})()
+    monkeypatch.setattr(safety, "devtools_endpoint", lambda ep: ("http://127.0.0.1:9222", None))
+    async def _mock_ws(ep):
+        return "ws://127.0.0.1:9222"
+    monkeypatch.setattr(browser_service, "devtools_ws_url", _mock_ws)
+
+    prof = BrowserProfileConfig(kind="attach", endpoint="http://127.0.0.1:9222")
+    await svc._attach(prof)
+
+    assert svc._attached is True
+    assert cdp_calls == [("Browser.setDownloadBehavior", {"behavior": "default", "eventsEnabled": True})]
+
+
+class _MockDownload:
+    def __init__(self, filename="report.pdf", url="https://example.com/report.pdf", page=None):
+        self.suggested_filename = filename
+        self.url = url
+        self.page = page
+        self.saved_to = None
+
+    async def save_as(self, path):
+        self.saved_to = path
+        Path(path).write_text("saved via download.save_as", encoding="utf-8")
+
+    async def cancel(self):
+        pass
+
+
+class _MockPageCDP:
+    def __init__(self, frame_id="sentient-frame"):
+        self.frame_id = frame_id
+        self.detached = False
+
+    async def send(self, command, params=None):
+        if command == "Page.getFrameTree":
+            return {"frameTree": {"frame": {"id": self.frame_id}, "childFrames": []}}
+        return {}
+
+    async def detach(self):
+        self.detached = True
+
+
+class _MockContext:
+    def __init__(self, cdp_session=None):
+        self.cdp_session = cdp_session
+
+    async def new_cdp_session(self, page):
+        if isinstance(self.cdp_session, Exception):
+            raise self.cdp_session
+        return self.cdp_session
+
+
+async def test_attached_download_matches_and_moves_valid_cdp_entry(tmp_path, monkeypatch):
+    monkeypatch.setattr(paths, "files_dir", lambda: tmp_path)
+    svc = BrowserService(make_app())
+    svc._attached = True
+    svc._attached_match_timeout_s = 0.05
+    page_cdp = _MockPageCDP("sentient-frame-1")
+    svc._context = _MockContext(page_cdp)
+
+    src_file = tmp_path / "browser-download.pdf"
+    src_file.write_text("content from browser", encoding="utf-8")
+
+    user_future = asyncio.get_running_loop().create_future()
+    user_entry = {
+        "guid": "g-user",
+        "frameId": "user-frame-99",
+        "filename": "report.pdf",
+        "url": "https://example.com/report.pdf",
+        "future": user_future,
+        "consumed": False,
+    }
+    sentient_future = asyncio.get_running_loop().create_future()
+    sentient_future.set_result(str(src_file))
+    sentient_entry = {
+        "guid": "g-sentient",
+        "frameId": "sentient-frame-1",
+        "filename": "report.pdf",
+        "url": "https://example.com/report.pdf",
+        "future": sentient_future,
+        "consumed": False,
+    }
+    svc._attached_downloads.extend([user_entry, sentient_entry])
+
+    rel = await svc._save_download(_MockDownload("report.pdf", page=object()))
+
+    assert rel == "downloads/report.pdf"
+    assert user_entry["consumed"] is False
+    assert not user_future.done()
+    assert sentient_entry["consumed"] is True
+    assert page_cdp.detached is True
+    saved_file = tmp_path / rel
+    assert saved_file.read_text(encoding="utf-8") == "content from browser"
+
+
+async def test_user_tab_entry_never_consumed_even_when_sentient_cdp_arrives_late(tmp_path, monkeypatch):
+    monkeypatch.setattr(paths, "files_dir", lambda: tmp_path)
+    svc = BrowserService(make_app())
+    svc._attached = True
+    svc._attached_match_timeout_s = 0.3
+    page_cdp = _MockPageCDP("sentient-frame-1")
+    svc._context = _MockContext(page_cdp)
+
+    src_file = tmp_path / "browser-download-late.pdf"
+    src_file.write_text("late content from browser", encoding="utf-8")
+
+    user_future = asyncio.get_running_loop().create_future()
+    user_entry = {
+        "guid": "g-user",
+        "frameId": "user-frame-99",
+        "filename": "report.pdf",
+        "url": "https://example.com/report.pdf",
+        "future": user_future,
+        "consumed": False,
+    }
+    svc._attached_downloads.append(user_entry)
+
+    task = asyncio.create_task(svc._save_download(_MockDownload("report.pdf", page=object())))
+    await asyncio.sleep(0.05)
+    assert user_entry["consumed"] is False
+
+    sentient_future = asyncio.get_running_loop().create_future()
+    sentient_future.set_result(str(src_file))
+    sentient_entry = {
+        "guid": "g-sentient",
+        "frameId": "sentient-frame-1",
+        "filename": "report.pdf",
+        "url": "https://example.com/report.pdf",
+        "future": sentient_future,
+        "consumed": False,
+    }
+    svc._attached_downloads.append(sentient_entry)
+
+    rel = await task
+    assert rel == "downloads/report.pdf"
+    assert user_entry["consumed"] is False
+    assert not user_future.done()
+    assert sentient_entry["consumed"] is True
+    assert (tmp_path / rel).read_text(encoding="utf-8") == "late content from browser"
+
+
+async def test_frame_tree_lookup_failure_uses_safe_fallback_without_consuming_user_entry(tmp_path, monkeypatch):
+    monkeypatch.setattr(paths, "files_dir", lambda: tmp_path)
+    svc = BrowserService(make_app())
+    svc._attached = True
+    svc._attached_match_timeout_s = 0.05
+    svc._context = _MockContext(RuntimeError("CDP detached"))
+
+    user_future = asyncio.get_running_loop().create_future()
+    user_entry = {
+        "guid": "g-user",
+        "frameId": "user-frame-99",
+        "filename": "report.pdf",
+        "url": "https://example.com/report.pdf",
+        "future": user_future,
+        "consumed": False,
+    }
+    svc._attached_downloads.append(user_entry)
+
+    dl = _MockDownload("report.pdf", page=object())
+    rel = await svc._save_download(dl)
+
+    assert rel == "downloads/report.pdf"
+    assert user_entry["consumed"] is False
+    assert not user_future.done()
+    assert dl.saved_to is not None
+    assert (tmp_path / rel).read_text(encoding="utf-8") == "saved via download.save_as"
+
+
+async def test_bounded_wait_terminates_when_no_matching_event_arrives(tmp_path, monkeypatch):
+    monkeypatch.setattr(paths, "files_dir", lambda: tmp_path)
+    svc = BrowserService(make_app())
+    svc._attached = True
+    svc._attached_match_timeout_s = 0.05
+    page_cdp = _MockPageCDP("sentient-frame-1")
+    svc._context = _MockContext(page_cdp)
+
+    user_future = asyncio.get_running_loop().create_future()
+    user_entry = {
+        "guid": "g-user",
+        "frameId": "user-frame-99",
+        "filename": "report.pdf",
+        "url": "https://example.com/report.pdf",
+        "future": user_future,
+        "consumed": False,
+    }
+    svc._attached_downloads.append(user_entry)
+
+    dl = _MockDownload("report.pdf", page=object())
+    t0 = time.monotonic()
+    rel = await svc._save_download(dl)
+    elapsed = time.monotonic() - t0
+
+    assert rel == "downloads/report.pdf"
+    assert user_entry["consumed"] is False
+    assert not user_future.done()
+    assert dl.saved_to is not None
+    assert elapsed >= 0.04
+    assert elapsed < 0.5
+
+
 def test_edge_downloads_hub_is_ignored_and_active_page_is_restored():
     svc = BrowserService(make_app())