diff --git a/src/browser_harness/helpers.py b/src/browser_harness/helpers.py index 3efb609c..bbffd6be 100644 --- a/src/browser_harness/helpers.py +++ b/src/browser_harness/helpers.py @@ -282,6 +282,9 @@ def capture_screenshot(path=None, full=False, max_dim=None): # --- tabs --- +_OPENED_TABS = set() +_KEEP_OPENED_TABS = False + def list_tabs(include_chrome=True): out = [] for t in cdp("Target.getTargets")["targetInfos"]: @@ -319,11 +322,55 @@ def new_tab(url="about:blank"): # attach, so the brief about:blank is "complete" by the time the caller # polls and wait_for_load() returns before navigation actually starts. tid = cdp("Target.createTarget", url="about:blank")["targetId"] + _OPENED_TABS.add(tid) switch_tab(tid) if url != "about:blank": goto_url(url) return tid +def opened_tabs(): + """Return targetIds opened by new_tab() in this CLI process.""" + return list(_OPENED_TABS) + +def keep_opened_tabs(keep=True): + """Opt out of automatic cleanup for tabs opened by this CLI process.""" + global _KEEP_OPENED_TABS + _KEEP_OPENED_TABS = keep + +def close_opened_tabs(force=False): + """Close tabs opened by new_tab() in this CLI process. + + browser-harness is commonly used from agents and cron jobs; leaving every + investigation tab open makes the visible profile unusable over time. This + helper intentionally only closes tabs created through new_tab() in the + current Python process, so pre-existing user/agent tabs are not touched. + """ + if _KEEP_OPENED_TABS and not force: + return [] + closed = [] + for tid in list(_OPENED_TABS): + try: + result = cdp("Target.closeTarget", targetId=tid) + if result.get("success", True): + closed.append(tid) + except Exception: + pass + finally: + _OPENED_TABS.discard(tid) + # Target.closeTarget is asynchronous in Chrome. Give it a short chance to + # settle so a follow-up browser-harness command does not see zombie tabs. + if closed: + deadline = time.time() + 2.0 + while time.time() < deadline: + try: + remaining = {t["targetId"] for t in list_tabs(include_chrome=True)} + except Exception: + break + if not any(tid in remaining for tid in closed): + break + time.sleep(0.05) + return closed + def ensure_real_tab(): """Switch to a real user tab if current is chrome:// / internal / stale.""" tabs = list_tabs(include_chrome=False) diff --git a/src/browser_harness/run.py b/src/browser_harness/run.py index c22ccab4..419a4360 100644 --- a/src/browser_harness/run.py +++ b/src/browser_harness/run.py @@ -103,7 +103,11 @@ def main(): ): start_remote_daemon(NAME) ensure_daemon() - exec(args[1], globals()) + try: + exec(args[1], globals()) + finally: + if os.environ.get("BH_KEEP_TABS") not in {"1", "true", "TRUE", "yes", "YES"}: + close_opened_tabs() if __name__ == "__main__": diff --git a/tests/unit/test_tab_cleanup.py b/tests/unit/test_tab_cleanup.py new file mode 100644 index 00000000..d918b93c --- /dev/null +++ b/tests/unit/test_tab_cleanup.py @@ -0,0 +1,94 @@ +import os +from unittest.mock import patch + +from browser_harness import helpers +from browser_harness import run + + +def test_new_tab_tracks_and_close_opened_tabs_only_closes_created_targets(): + helpers._OPENED_TABS.clear() + helpers.keep_opened_tabs(False) + calls = [] + + def fake_cdp(method, **kwargs): + calls.append((method, kwargs)) + if method == "Target.createTarget": + return {"targetId": "tab-created"} + if method == "Target.attachToTarget": + return {"sessionId": "session-created"} + return {} + + with patch("browser_harness.helpers.cdp", side_effect=fake_cdp), \ + patch("browser_harness.helpers.goto_url", return_value={}): + assert helpers.new_tab("https://example.com") == "tab-created" + assert helpers.opened_tabs() == ["tab-created"] + assert helpers.close_opened_tabs() == ["tab-created"] + + assert ("Target.closeTarget", {"targetId": "tab-created"}) in calls + assert helpers.opened_tabs() == [] + + +def test_keep_opened_tabs_opt_out_until_forced(): + helpers._OPENED_TABS.clear() + helpers.keep_opened_tabs(True) + + with patch("browser_harness.helpers.cdp") as mock_cdp: + helpers._OPENED_TABS.add("tab-keep") + assert helpers.close_opened_tabs() == [] + mock_cdp.assert_not_called() + assert helpers.close_opened_tabs(force=True) == ["tab-keep"] + + helpers.keep_opened_tabs(False) + + +def test_cli_auto_closes_opened_tabs_in_finally(monkeypatch): + monkeypatch.delenv("BH_KEEP_TABS", raising=False) + monkeypatch.setattr(run.sys, "argv", ["browser-harness", "-c", "new_tab('https://example.com')\nraise RuntimeError('boom')"]) + events = [] + + def fake_cdp(method, **kwargs): + events.append((method, kwargs)) + if method == "Target.createTarget": + return {"targetId": "tab-cli"} + if method == "Target.attachToTarget": + return {"sessionId": "session-cli"} + return {} + + with patch("browser_harness.run.print_update_banner"), \ + patch("browser_harness.run.daemon_alive", return_value=True), \ + patch("browser_harness.run.ensure_daemon"), \ + patch("browser_harness.helpers.cdp", side_effect=fake_cdp), \ + patch("browser_harness.helpers.goto_url", return_value={}): + try: + run.main() + except RuntimeError as exc: + assert str(exc) == "boom" + else: + raise AssertionError("RuntimeError was not raised") + + assert ("Target.closeTarget", {"targetId": "tab-cli"}) in events + + +def test_cli_respects_bh_keep_tabs(monkeypatch): + monkeypatch.setenv("BH_KEEP_TABS", "1") + monkeypatch.setattr(run.sys, "argv", ["browser-harness", "-c", "new_tab('https://example.com')"]) + events = [] + + def fake_cdp(method, **kwargs): + events.append((method, kwargs)) + if method == "Target.createTarget": + return {"targetId": "tab-kept-by-env"} + if method == "Target.attachToTarget": + return {"sessionId": "session-kept"} + return {} + + with patch("browser_harness.run.print_update_banner"), \ + patch("browser_harness.run.daemon_alive", return_value=True), \ + patch("browser_harness.run.ensure_daemon"), \ + patch("browser_harness.helpers.cdp", side_effect=fake_cdp), \ + patch("browser_harness.helpers.goto_url", return_value={}): + run.main() + + assert not any(method == "Target.closeTarget" for method, _ in events) + helpers._OPENED_TABS.clear() + os.environ.pop("BH_KEEP_TABS", None)