From fd1bb2f8785a20c7fe5da914fbdae9840977b792 Mon Sep 17 00:00:00 2001
From: joaoneves0 <joaoneves@reps2.com>
Date: Mon, 14 Apr 2025 14:23:38 -0300
Subject: [PATCH] Ensure virtual display cleanup runs on browser close failure

---
 pythonlib/camoufox/utils.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/pythonlib/camoufox/utils.py b/pythonlib/camoufox/utils.py
index 7158ba8..0c255b3 100644
--- a/pythonlib/camoufox/utils.py
+++ b/pythonlib/camoufox/utils.py
@@ -302,9 +302,13 @@ async def async_attach_vd(
     _close = browser.close
 
     async def new_close(*args: Any, **kwargs: Any):
-        await _close(*args, **kwargs)
-        if virtual_display:
-            virtual_display.kill()
+        try:
+            await _close(*args, **kwargs)
+        except Exception:
+            raise
+        finally:
+            if virtual_display:
+                virtual_display.kill()
 
     browser.close = new_close
     browser._virtual_display = virtual_display
@@ -324,9 +328,13 @@ def sync_attach_vd(
     _close = browser.close
 
     def new_close(*args: Any, **kwargs: Any):
-        _close(*args, **kwargs)
-        if virtual_display:
-            virtual_display.kill()
+        try:
+            _close(*args, **kwargs)
+        except Exception:
+            raise
+        finally:
+            if virtual_display:
+                virtual_display.kill()
 
     browser.close = new_close
     browser._virtual_display = virtual_display