Browser reopened if closed (#314)

mentat/session.py

@@ -180,7 +180,9 @@ async def _stop(self):
 
         session_context = SESSION_CONTEXT.get()
         cost_tracker = session_context.cost_tracker
+        vision_manager = session_context.vision_manager
 
+        vision_manager.close()
         cost_tracker.display_total_cost()
         logging.shutdown()
         self._exit_task.cancel()

mentat/vision/vision_manager.py

@@ -4,6 +4,7 @@
 
 import attr
 from selenium import webdriver
+from selenium.common.exceptions import NoSuchWindowException
 from selenium.webdriver.chrome.service import Service
 from webdriver_manager.chrome import ChromeDriverManager
 
@@ -19,14 +20,22 @@ class VisionManager:
     driver: webdriver.Chrome | None = attr.field(default=None)
 
     def _open_browser(self) -> None:
-        if self.driver is None:
+        if self.driver is None or not self.driver_running():
             service = Service(ChromeDriverManager().install())
             self.driver = webdriver.Chrome(service=service)
 
     def open(self, path: str) -> None:
         self._open_browser()
         self.driver.get(path)  # type: ignore
 
+    def driver_running(self):
+        try:
+            # This command should fail if the driver is not running
+            self.driver.execute_script('return "hello world";')  # type: ignore
+            return True
+        except NoSuchWindowException:
+            return False
+
     def screenshot(self, path: Optional[str] = None) -> str:
         if path is not None:
             expanded = os.path.abspath(os.path.expanduser(path))
@@ -46,3 +55,7 @@ def screenshot(self, path: Optional[str] = None) -> str:
         image_data = f"data:image/png;base64,{decoded}"
 
         return image_data
+
+    def close(self) -> None:
+        if self.driver is not None:
+            self.driver.quit()