1 files changed, 14 insertions, 20 deletions

diff --git a/g4f/Provider/MyShell.py b/g4f/Provider/MyShell.py
index 548c4be1..a1c8d335 100644
--- a/g4f/Provider/MyShell.py
+++ b/g4f/Provider/MyShell.py
@@ -4,7 +4,7 @@ import time, json
 
 from ..typing import CreateResult, Messages
 from .base_provider import BaseProvider
-from .helper import WebDriver, format_prompt, get_browser
+from .helper import WebDriver, WebDriverSession, format_prompt
 
 class MyShell(BaseProvider):
     url = "https://app.myshell.ai/chat"
@@ -20,22 +20,27 @@ class MyShell(BaseProvider):
         stream: bool,
         proxy: str = None,
         timeout: int = 120,
-        browser: WebDriver = None,
+        web_driver: WebDriver = None,
         **kwargs
     ) -> CreateResult:
-        driver = browser if browser else get_browser("", False, proxy)
+        with WebDriverSession(web_driver, "", proxy=proxy) as driver:
+            from selenium.webdriver.common.by import By
+            from selenium.webdriver.support.ui import WebDriverWait
+            from selenium.webdriver.support import expected_conditions as EC
 
-        from selenium.webdriver.common.by import By
-        from selenium.webdriver.support.ui import WebDriverWait
-        from selenium.webdriver.support import expected_conditions as EC
+            driver.get(cls.url)
 
-        driver.get(cls.url)
-        try:
             # Wait for page load and cloudflare validation
             WebDriverWait(driver, timeout).until(
                 EC.presence_of_element_located((By.CSS_SELECTOR, "body:not(.no-js)"))
             )
             # Send request with message
+            data = {
+                "botId": "4738",
+                "conversation_scenario": 3,
+                "message": format_prompt(messages),
+                "messageType": 1
+            }
             script = """
 response = await fetch("https://api.myshell.ai/v1/bot/chat/send_message", {
     "headers": {
@@ -49,12 +54,6 @@ response = await fetch("https://api.myshell.ai/v1/bot/chat/send_message", {
 })
 window.reader = response.body.getReader();
 """
-            data = {
-                "botId": "4738",
-                "conversation_scenario": 3,
-                "message": format_prompt(messages),
-                "messageType": 1
-            }
             driver.execute_script(script.replace("{body}", json.dumps(data)))
             script = """
 chunk = await window.reader.read();
@@ -80,9 +79,4 @@ return content;
                 elif chunk != "":
                     break
                 else:
-                    time.sleep(0.1)
-        finally:
-            if not browser:
-                driver.close()
-                time.sleep(0.1)
-                driver.quit()
\ No newline at end of file
+                    time.sleep(0.1)
\ No newline at end of file