diff --git a/backend/packages/harness/deerflow/community/browser_automation/session.py b/backend/packages/harness/deerflow/community/browser_automation/session.py index 36353a3b6..148b2052e 100644 --- a/backend/packages/harness/deerflow/community/browser_automation/session.py +++ b/backend/packages/harness/deerflow/community/browser_automation/session.py @@ -21,6 +21,7 @@ import os import threading import time from collections.abc import Callable, Coroutine +from concurrent.futures import Future from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, TypeVar from urllib.parse import urlparse @@ -172,6 +173,14 @@ def _is_playwright_timeout_error(exc: Exception) -> bool: return exc.__class__.__name__ == "TimeoutError" and exc.__class__.__module__.startswith("playwright.") +def _consume_future_exception(future: asyncio.Future[Any]) -> None: + """Retrieve detached close errors; the concurrent-future callback logs them.""" + if future.cancelled(): + return + with contextlib.suppress(Exception): + future.exception() + + def redact_browser_url(url: str) -> str: """Drop query/fragment so a blocked-URL log line can't leak tokens/PII.""" try: @@ -198,7 +207,7 @@ class _PlaywrightLoopThread: future = asyncio.run_coroutine_threadsafe(coro, self._loop) return await asyncio.wrap_future(future) - def submit(self, coro: Coroutine[Any, Any, Any]) -> None: + def submit(self, coro: Coroutine[Any, Any, Any]) -> Future[Any]: """Schedule *coro* on the private loop without blocking the caller.""" future = asyncio.run_coroutine_threadsafe(coro, self._loop) @@ -209,6 +218,7 @@ class _PlaywrightLoopThread: logger.debug("browser background task failed: %s", exc) future.add_done_callback(_log_failure) + return future def run_sync(self, coro: Coroutine[Any, Any, T], timeout: float | None = None) -> T: future = asyncio.run_coroutine_threadsafe(coro, self._loop) @@ -849,8 +859,18 @@ class BrowserSession: with self._activity(): await self._loop.run(self._dispatch_input(event)) + def _submit_close(self) -> Future[Any]: + close_coro = self._close() + try: + return self._loop.submit(close_coro) + except Exception: + close_coro.close() + raise + async def close(self) -> None: - await self._loop.run(self._close()) + close_future = asyncio.wrap_future(self._submit_close()) + close_future.add_done_callback(_consume_future_exception) + await asyncio.shield(close_future) class BrowserSessionManager: @@ -1019,8 +1039,19 @@ class BrowserSessionManager: sessions = list(self._sessions.values()) self._sessions.clear() self._last_used.clear() + close_futures: list[asyncio.Future[Any]] = [] for session in sessions: - await session.close() + try: + close_future = asyncio.wrap_future(session._submit_close()) + except Exception as exc: + logger.debug("browser session close submission failed: %s", exc) + continue + close_future.add_done_callback(_consume_future_exception) + close_futures.append(close_future) + if close_futures: + close_group = asyncio.gather(*close_futures) + close_group.add_done_callback(_consume_future_exception) + await asyncio.shield(close_group) return len(sessions) diff --git a/backend/tests/test_browser_close_cancellation.py b/backend/tests/test_browser_close_cancellation.py new file mode 100644 index 000000000..75738d72c --- /dev/null +++ b/backend/tests/test_browser_close_cancellation.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import asyncio +from concurrent.futures import Future +from typing import Any, cast + +import pytest + +import deerflow.community.browser_automation.session as session_module +from deerflow.community.browser_automation.session import BrowserSession, BrowserSessionManager + + +class _ControllablePrivateLoop: + """Model the cancellation boundary between a caller loop and Playwright's loop.""" + + def __init__(self) -> None: + self.cleanup_futures: list[Future[None]] = [] + self.run_cancelled = False + + async def run(self, coro: Any) -> None: + # Current main awaits the private-loop proxy directly, so caller + # cancellation propagates through wrap_future and cancels cleanup. + coro.close() + try: + await asyncio.Future() + except asyncio.CancelledError: + self.run_cancelled = True + raise + + def submit(self, coro: Any) -> Future[None]: + # The fixed close path hands cleanup to the private loop and awaits its + # concurrent future behind a shield. Closing the coroutine here avoids + # needing Playwright in this focused lifecycle test. + coro.close() + cleanup_future: Future[None] = Future() + self.cleanup_futures.append(cleanup_future) + return cleanup_future + + +class _FailingPrivateLoop(_ControllablePrivateLoop): + def __init__(self) -> None: + super().__init__() + self.submitted_coro: Any | None = None + + def submit(self, coro: Any) -> Future[None]: + self.submitted_coro = coro + raise RuntimeError("private loop closed") + + +def _session(loop: _ControllablePrivateLoop) -> BrowserSession: + return BrowserSession( + cast(Any, loop), + headless=True, + timeout_ms=1000, + viewport={"width": 1000, "height": 500}, + ) + + +@pytest.mark.asyncio +async def test_browser_close_caller_cancellation_does_not_cancel_private_cleanup() -> None: + loop = _ControllablePrivateLoop() + session = _session(loop) + + close_task = asyncio.create_task(session.close()) + await asyncio.sleep(0) + + close_task.cancel("caller stopped") + with pytest.raises(asyncio.CancelledError): + await close_task + + assert not loop.run_cancelled + assert len(loop.cleanup_futures) == 1 + assert not loop.cleanup_futures[0].cancelled() + + loop.cleanup_futures[0].set_result(None) + await asyncio.sleep(0) + + +def test_submit_close_closes_coroutine_when_submission_fails() -> None: + loop = _FailingPrivateLoop() + session = _session(loop) + + with pytest.raises(RuntimeError, match="private loop closed"): + session._submit_close() + + assert loop.submitted_coro is not None + assert loop.submitted_coro.cr_frame is None + + +@pytest.mark.asyncio +async def test_manager_close_session_cancellation_keeps_detached_cleanup_running() -> None: + loop = _ControllablePrivateLoop() + session = _session(loop) + manager = BrowserSessionManager() + manager._sessions["thread-a"] = session + manager._last_used["thread-a"] = 0.0 + + close_task = asyncio.create_task(manager.close_session("thread-a")) + await asyncio.sleep(0) + + assert "thread-a" not in manager._sessions + assert len(loop.cleanup_futures) == 1 + + close_task.cancel("request stopped") + with pytest.raises(asyncio.CancelledError): + await close_task + + assert not loop.cleanup_futures[0].cancelled() + loop.cleanup_futures[0].set_result(None) + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_manager_close_all_submits_every_cleanup_before_cancellable_wait() -> None: + loop = _ControllablePrivateLoop() + manager = BrowserSessionManager() + manager._sessions.update({"thread-a": _session(loop), "thread-b": _session(loop)}) + manager._last_used.update({"thread-a": 0.0, "thread-b": 0.0}) + + close_task = asyncio.create_task(manager.close_all_sessions()) + await asyncio.sleep(0) + + assert manager._sessions == {} + assert len(loop.cleanup_futures) == 2 + + close_task.cancel("shutdown interrupted") + with pytest.raises(asyncio.CancelledError): + await close_task + + assert all(not future.cancelled() for future in loop.cleanup_futures) + for future in loop.cleanup_futures: + future.set_result(None) + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_manager_close_all_continues_after_sync_submission_failure() -> None: + failing_loop = _FailingPrivateLoop() + healthy_loop = _ControllablePrivateLoop() + manager = BrowserSessionManager() + manager._sessions.update({"broken": _session(failing_loop), "healthy": _session(healthy_loop)}) + manager._last_used.update({"broken": 0.0, "healthy": 0.0}) + + close_task = asyncio.create_task(manager.close_all_sessions()) + await asyncio.sleep(0) + + assert manager._sessions == {} + assert len(healthy_loop.cleanup_futures) == 1 + assert failing_loop.submitted_coro is not None + assert failing_loop.submitted_coro.cr_frame is None + + healthy_loop.cleanup_futures[0].set_result(None) + assert await close_task == 2 + + +@pytest.mark.asyncio +async def test_manager_close_all_consumes_group_failure_after_caller_cancellation(monkeypatch: pytest.MonkeyPatch) -> None: + loop = _ControllablePrivateLoop() + manager = BrowserSessionManager() + manager._sessions["thread-a"] = _session(loop) + manager._last_used["thread-a"] = 0.0 + + original_gather = asyncio.gather + original_consume = session_module._consume_future_exception + group_future: asyncio.Future[Any] | None = None + consumed: list[asyncio.Future[Any]] = [] + + def tracking_gather(*aws: Any, **kwargs: Any) -> asyncio.Future[Any]: + nonlocal group_future + group_future = cast(asyncio.Future[Any], original_gather(*aws, **kwargs)) + return group_future + + def recording_consume(future: asyncio.Future[Any]) -> None: + consumed.append(future) + original_consume(future) + + monkeypatch.setattr(asyncio, "gather", tracking_gather) + monkeypatch.setattr(session_module, "_consume_future_exception", recording_consume) + + close_task = asyncio.create_task(manager.close_all_sessions()) + await asyncio.sleep(0) + assert group_future is not None + + close_task.cancel("shutdown interrupted") + with pytest.raises(asyncio.CancelledError): + await close_task + + loop.cleanup_futures[0].set_exception(RuntimeError("teardown failed")) + await asyncio.sleep(0) + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert group_future.done() + assert group_future in consumed