1 files changed, 49 insertions, 43 deletions

diff --git a/g4f/Provider/needs_auth/Theb.py b/g4f/Provider/needs_auth/Theb.py
index 89c69727..cf33f0c6 100644
--- a/g4f/Provider/needs_auth/Theb.py
+++ b/g4f/Provider/needs_auth/Theb.py
@@ -4,7 +4,7 @@ import time
 
 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
 
 models = {
     "theb-ai": "TheB.AI",
@@ -44,26 +44,60 @@ class Theb(BaseProvider):
         messages: Messages,
         stream: bool,
         proxy: str = None,
-        browser: WebDriver = None,
-        headless: bool = True,
+        web_driver: WebDriver = None,
+        virtual_display: bool = True,
         **kwargs
     ) -> CreateResult:
         if model in models:
             model = models[model]
         prompt = format_prompt(messages)
-        driver = browser if browser else get_browser(None, headless, proxy)
+        web_session = WebDriverSession(web_driver, virtual_display=virtual_display, proxy=proxy)
+        with web_session 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.keys import Keys
 
-        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.keys import Keys
+            # Register fetch hook
+            script = """
+window._fetch = window.fetch;
+window.fetch = (url, options) => {
+    // Call parent fetch method
+    const result = window._fetch(url, options);
+    if (!url.startsWith("/api/conversation")) {
+        return result;
+    }
+    // Load response reader
+    result.then((response) => {
+        if (!response.body.locked) {
+            window._reader = response.body.getReader();
+        }
+    });
+    // Return dummy response
+    return new Promise((resolve, reject) => {
+        resolve(new Response(new ReadableStream()))
+    });
+}
+window._last_message = "";
+"""
+            driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", {
+                "source": script
+            })
+
+            try:
+                driver.get(f"{cls.url}/home")
+                wait = WebDriverWait(driver, 5)
+                wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize")))
+            except:
+                driver = web_session.reopen()
+                driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", {
+                    "source": script
+                })
+                driver.get(f"{cls.url}/home")
+                wait = WebDriverWait(driver, 240)
+                wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize")))
 
-        
-        try:
-            driver.get(f"{cls.url}/home")
-            wait = WebDriverWait(driver, 10 if headless else 240)
-            wait.until(EC.visibility_of_element_located((By.TAG_NAME, "body")))
-            time.sleep(0.1)
+            time.sleep(200)
             try:
                 driver.find_element(By.CSS_SELECTOR, ".driver-overlay").click()
                 driver.find_element(By.CSS_SELECTOR, ".driver-overlay").click()
@@ -87,29 +121,6 @@ class Theb(BaseProvider):
                 button = container.find_element(By.CSS_SELECTOR, "button.btn-blue.btn-small.border")
                 button.click()
 
-            # Register fetch hook
-            script = """
-window._fetch = window.fetch;
-window.fetch = (url, options) => {
-    // Call parent fetch method
-    const result = window._fetch(url, options);
-    if (!url.startsWith("/api/conversation")) {
-        return result;
-    }
-    // Load response reader
-    result.then((response) => {
-        if (!response.body.locked) {
-            window._reader = response.body.getReader();
-        }
-    });
-    // Return dummy response
-    return new Promise((resolve, reject) => {
-        resolve(new Response(new ReadableStream()))
-    });
-}
-window._last_message = "";
-"""
-            driver.execute_script(script)
 
             # Submit prompt
             wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize")))
@@ -150,9 +161,4 @@ return '';
                 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