Update app/lens_text_core.py

app/lens_text_core.py

@@ -1,7 +1,7 @@
 
 import os, time, asyncio, base64, re, threading, hashlib, logging, shutil, atexit, tempfile, uuid
 from io import BytesIO
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Optional
 from urllib.parse import urlparse
 
 import httpx
@@ -33,55 +33,83 @@ CHROME_EXTRA_ARGS = os.getenv(
     "--window-size=1920,1080 --headless=new",
 ).split()
 
-_PROFILE_DIRS = []
-
 _CACHE_TTL    = 300  
 _IDLE_TIMEOUT = int(os.getenv("CHROME_IDLE_SECONDS", "60"))
 _BROWSER_TTL  = 900  
 
 def _mk_profile_dir() -> str:
-    base = os.getenv("CHROME_PROFILE_BASE", tempfile.gettempdir())
+    base = tempfile.gettempdir()
     p = os.path.join(base, f"chrome-profile-{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}")
     os.makedirs(p, exist_ok=True)
-    _PROFILE_DIRS.append(p)
+    
+    atexit.register(lambda: shutil.rmtree(p, ignore_errors=True))
     return p
 
-def _cleanup_profiles():
-    for p in _PROFILE_DIRS:
-        try:
-            shutil.rmtree(p, ignore_errors=True)
-        except Exception:
-            pass
-
-atexit.register(_cleanup_profiles)
-
-def _build_chrome() -> webdriver.Chrome:
+def _build_chrome(cookie_dict: Dict[str, str] | None = None,
+                  *, domain: str = ".google.com",
+                  user_data_dir: Optional[str] = None,
+                  reuse_profile: bool = False):
     opts = ChromeOptions()
 
     extra = os.getenv(
         "CHROME_EXTRA_ARGS",
-        "--disable-gpu --no-sandbox --disable-dev-shm-usage --window-size=1920,1080 --headless=new",
+        "--no-sandbox --disable-dev-shm-usage --disable-gpu --window-size=1920,1080 --headless=new",
     ).split()
     for a in extra:
         if a:
             opts.add_argument(a)
 
-    profile_dir = _mk_profile_dir()
+    if user_data_dir:
+        profile_dir = user_data_dir
+        os.makedirs(profile_dir, exist_ok=True)
+    elif reuse_profile:
+        base = os.getenv("CHROME_PROFILE_BASE", "/tmp")
+        profile_dir = os.path.join(base, "chrome-profile")
+        os.makedirs(profile_dir, exist_ok=True)
+    else:
+        profile_dir = _mk_profile_dir()
+
     opts.add_argument(f"--user-data-dir={profile_dir}")
     opts.add_argument("--profile-directory=Default")
 
+    bin_loc = os.getenv("CHROME_BINARY")
+    if bin_loc and os.path.exists(bin_loc):
+        opts.binary_location = bin_loc
+
     drv_path = os.getenv("CHROMEDRIVER")
-    try:
+
+    def _new_driver():
         if drv_path and os.path.exists(drv_path):
-            return webdriver.Chrome(service=ChromeService(executable_path=drv_path), options=opts)
-        return webdriver.Chrome(options=opts)
-    except SessionNotCreatedException as e:
-        LOGGER.warning("SessionNotCreated: %s; retry with a fresh profile dir", e)
-        profile_dir2 = _mk_profile_dir()
-        opts.arguments = [a for a in opts.arguments if not a.startswith("--user-data-dir=")]
-        opts.add_argument(f"--user-data-dir={profile_dir2}")
+            service = ChromeService(executable_path=drv_path)
+            return webdriver.Chrome(service=service, options=opts)
         return webdriver.Chrome(options=opts)
 
+    try:
+        drv = _new_driver()
+    except SessionNotCreatedException:
+        fresh = _mk_profile_dir()
+        opts.arguments = [x for x in opts.arguments if not x.startswith("--user-data-dir=")]
+        opts.add_argument(f"--user-data-dir={fresh}")
+        drv = _new_driver()
+
+    if cookie_dict:
+        try:
+            drv.execute_cdp_cmd("Network.enable", {})
+            cookies = [
+                {"name": k, "value": v, "domain": domain, "path": "/", "secure": True}
+                for k, v in cookie_dict.items()
+            ]
+            drv.execute_cdp_cmd("Network.setCookies", {"cookies": cookies})
+        except Exception:
+            try:
+                drv.get(f"https://{domain.lstrip('.')}/")
+                for k, v in cookie_dict.items():
+                    drv.add_cookie({"name": k, "value": v, "domain": domain, "path": "/", "secure": True})
+            except Exception:
+                pass
+
+    return drv
+
 _cached_cookie, _cached_cookie_ts, _cookie_lock = None, 0.0, threading.Lock()
 _global_driver, _driver_last_use, _driver_lock  = None, 0.0, threading.Lock()