refactor(setup): streamline Chrome driver creation and options

Author: stickerdaniel

linkedin_mcp_server/drivers/chrome.py

@@ -8,8 +8,6 @@
 
 import logging
 import os
-import shutil
-import tempfile
 from typing import Dict, Optional
 
 from linkedin_scraper.exceptions import (
@@ -34,29 +32,22 @@
 # Global driver storage to reuse sessions
 active_drivers: Dict[str, webdriver.Chrome] = {}
 
-# Store user data directories for cleanup
-user_data_dirs: Dict[str, str] = {}
 
 logger = logging.getLogger(__name__)
 
 
-def create_chrome_driver(session_id: str = "default") -> webdriver.Chrome:
+def create_chrome_options(config) -> Options:
     """
-    Create a new Chrome WebDriver instance with proper configuration.
+    Create Chrome options with all necessary configuration for LinkedIn scraping.
 
     Args:
-        session_id: Unique identifier for the session (used for cleanup)
+        config: AppConfig instance with Chrome configuration
 
     Returns:
-        webdriver.Chrome: Configured Chrome WebDriver instance
-
-    Raises:
-        WebDriverException: If driver creation fails
+        Options: Configured Chrome options object
     """
-    config = get_config()
-
-    # Set up Chrome options
     chrome_options = Options()
+
     logger.info(
         f"Running browser in {'headless' if config.chrome.headless else 'visible'} mode"
     )
@@ -70,14 +61,15 @@ def create_chrome_driver(session_id: str = "default") -> webdriver.Chrome:
     chrome_options.add_argument("--window-size=1920,1080")
     chrome_options.add_argument("--disable-extensions")
     chrome_options.add_argument("--disable-background-timer-throttling")
-
-    # Create a unique user data directory to avoid conflicts
-    user_data_dir = tempfile.mkdtemp(prefix="linkedin_mcp_chrome_")
-    chrome_options.add_argument(f"--user-data-dir={user_data_dir}")
-    logger.debug(f"Using Chrome user data directory: {user_data_dir}")
-
-    # Store the user data directory for cleanup
-    user_data_dirs[session_id] = user_data_dir
+    chrome_options.add_argument("--disable-background-networking")
+    chrome_options.add_argument("--disable-default-apps")
+    chrome_options.add_argument("--disable-sync")
+    chrome_options.add_argument("--metrics-recording-only")
+    chrome_options.add_argument("--no-default-browser-check")
+    chrome_options.add_argument("--no-first-run")
+    chrome_options.add_argument("--disable-features=TranslateUI,BlinkGenPropertyTrees")
+    chrome_options.add_argument("--aggressive-cache-discard")
+    chrome_options.add_argument("--disable-ipc-flooding-protection")
 
     # Set user agent (configurable with sensible default)
     user_agent = getattr(config.chrome, "user_agent", DEFAULT_USER_AGENT)
@@ -87,20 +79,96 @@ def create_chrome_driver(session_id: str = "default") -> webdriver.Chrome:
     for arg in config.chrome.browser_args:
         chrome_options.add_argument(arg)
 
-    # Initialize Chrome driver
-    logger.info("Initializing Chrome WebDriver...")
+    return chrome_options
 
+
+def create_chrome_service(config):
+    """
+    Create Chrome service with ChromeDriver path resolution.
+
+    Args:
+        config: AppConfig instance with Chrome configuration
+
+    Returns:
+        Service or None: Chrome service if path is configured, None for auto-detection
+    """
     # Use ChromeDriver path from environment or config
     chromedriver_path = (
         os.environ.get("CHROMEDRIVER_PATH") or config.chrome.chromedriver_path
     )
 
     if chromedriver_path:
         logger.info(f"Using ChromeDriver at path: {chromedriver_path}")
-        service = Service(executable_path=chromedriver_path)
-        driver = webdriver.Chrome(service=service, options=chrome_options)
+        return Service(executable_path=chromedriver_path)
     else:
         logger.info("Using auto-detected ChromeDriver")
+        return None
+
+
+def create_temporary_chrome_driver() -> webdriver.Chrome:
+    """
+    Create a temporary Chrome WebDriver instance for one-off operations.
+
+    This driver is NOT stored in the global active_drivers dict and should be
+    manually cleaned up by the caller.
+
+    Returns:
+        webdriver.Chrome: Configured Chrome WebDriver instance
+
+    Raises:
+        WebDriverException: If driver creation fails
+    """
+    config = get_config()
+
+    logger.info("Creating temporary Chrome WebDriver...")
+
+    # Create Chrome options using shared function
+    chrome_options = create_chrome_options(config)
+
+    # Create Chrome service using shared function
+    service = create_chrome_service(config)
+
+    # Initialize Chrome driver
+    if service:
+        driver = webdriver.Chrome(service=service, options=chrome_options)
+    else:
+        driver = webdriver.Chrome(options=chrome_options)
+
+    logger.info("Temporary Chrome WebDriver created successfully")
+
+    # Add a page load timeout for safety
+    driver.set_page_load_timeout(60)
+
+    # Set shorter implicit wait for faster operations
+    driver.implicitly_wait(10)
+
+    return driver
+
+
+def create_chrome_driver() -> webdriver.Chrome:
+    """
+    Create a new Chrome WebDriver instance with proper configuration.
+
+    Returns:
+        webdriver.Chrome: Configured Chrome WebDriver instance
+
+    Raises:
+        WebDriverException: If driver creation fails
+    """
+    config = get_config()
+
+    logger.info("Initializing Chrome WebDriver...")
+
+    # Create Chrome options using shared function
+    chrome_options = create_chrome_options(config)
+
+    # Create Chrome service using shared function
+    service = create_chrome_service(config)
+
+    # Initialize Chrome driver
+    if service:
+        driver = webdriver.Chrome(service=service, options=chrome_options)
+    else:
         driver = webdriver.Chrome(options=chrome_options)
 
     logger.info("Chrome WebDriver initialized successfully")
@@ -229,7 +297,7 @@ def get_or_create_driver(authentication: str) -> webdriver.Chrome:
 
     try:
         # Create new driver
-        driver = create_chrome_driver(session_id)
+        driver = create_chrome_driver()
 
         # Login to LinkedIn
         login_to_linkedin(driver, authentication)
@@ -261,7 +329,7 @@ def get_or_create_driver(authentication: str) -> webdriver.Chrome:
 
 def close_all_drivers() -> None:
     """Close all active drivers and clean up resources."""
-    global active_drivers, user_data_dirs
+    global active_drivers
 
     for session_id, driver in active_drivers.items():
         try:
@@ -270,21 +338,8 @@ def close_all_drivers() -> None:
         except Exception as e:
             logger.warning(f"Error closing driver {session_id}: {e}")
 
-        # Clean up user data directory
-        if session_id in user_data_dirs:
-            try:
-                user_data_dir = user_data_dirs[session_id]
-                if os.path.exists(user_data_dir):
-                    shutil.rmtree(user_data_dir)
-                    logger.debug(f"Cleaned up user data directory: {user_data_dir}")
-            except Exception as e:
-                logger.warning(
-                    f"Error cleaning up user data directory for session {session_id}: {e}"
-                )
-
     active_drivers.clear()
-    user_data_dirs.clear()
-    logger.info("All Chrome WebDriver sessions closed and cleaned up")
+    logger.info("All Chrome WebDriver sessions closed")
 
 
 def get_active_driver() -> Optional[webdriver.Chrome]:

linkedin_mcp_server/setup.py

@@ -6,14 +6,11 @@
 """
 
 import logging
-import os
-import tempfile
 from contextlib import contextmanager
 from typing import Dict, Iterator
 
 import inquirer
 from selenium import webdriver
-from selenium.webdriver.chrome.options import Options
 
 from linkedin_mcp_server.authentication import store_authentication
 from linkedin_mcp_server.config import get_config
@@ -105,44 +102,13 @@ def temporary_chrome_driver() -> Iterator[webdriver.Chrome]:
     Raises:
         Exception: If driver creation fails
     """
-    config: AppConfig = get_config()
-
-    logger.info("Creating temporary browser for cookie capture...")
-
-    # Set up Chrome options for cookie capture
-    chrome_options = Options()
-    if config.chrome.headless:
-        chrome_options.add_argument("--headless=new")
-
-    # Add essential options
-    # chrome_options.add_argument("--no-sandbox")
-    # chrome_options.add_argument("--disable-dev-shm-usage")
-    # chrome_options.add_argument("--disable-gpu")
-    # chrome_options.add_argument("--window-size=3456,2234")
-
-    # Create a unique user data directory to avoid conflicts
-    user_data_dir = tempfile.mkdtemp(prefix="linkedin_mcp_setup_")
-    chrome_options.add_argument(f"--user-data-dir={user_data_dir}")
-    logger.debug(f"Using Chrome user data directory for setup: {user_data_dir}")
+    from linkedin_mcp_server.drivers.chrome import create_temporary_chrome_driver
 
     driver = None
     try:
-        # Create temporary driver
-        chromedriver_path = (
-            os.environ.get("CHROMEDRIVER_PATH") or config.chrome.chromedriver_path
-        )
-
-        if chromedriver_path:
-            from selenium.webdriver.chrome.service import Service
-
-            service = Service(executable_path=chromedriver_path)
-            driver = webdriver.Chrome(service=service, options=chrome_options)
-        else:
-            driver = webdriver.Chrome(options=chrome_options)
-
-        driver.set_page_load_timeout(60)
+        # Create temporary driver using shared function
+        driver = create_temporary_chrome_driver()
         yield driver
-
     finally:
         if driver:
             driver.quit()
@@ -174,7 +140,7 @@ def capture_cookie_from_credentials(email: str, password: str) -> str:
             email,
             password,
             timeout=60,  # longer timeout for login (captcha, mobile verification, etc.)
-            interactive=interactive,  # Respect configuration setting
+            interactive=interactive,  # type: ignore  # Respect configuration setting
         )
 
         # Capture cookie

main.py

@@ -80,16 +80,20 @@ def get_cookie_and_exit() -> None:
         print(cookie)
 
         # Try to copy to clipboard
+        clipboard_success = False
         try:
             import pyperclip
 
             pyperclip.copy(cookie)
+            clipboard_success = True
+            print("📋 Cookie copied to clipboard!")
+        except Exception as e:
+            logger.debug(f"pyperclip clipboard failed: {e}")
+
+        if not clipboard_success:
             print(
-                "📋 Cookie copied to clipboard! Now you can set the LINKEDIN_COOKIE environment variable in your configuration"
+                "💡 Set this cookie as an environment variable in your config or pass it with --cookie flag"
             )
-        except Exception as e:
-            logger.warning(f"Could not copy to clipboard: {e}")
-            print("⚠️  Copy the cookie above manually")
 
     except Exception as e:
         logger.error(f"Error getting cookie: {e}")
@@ -108,7 +112,7 @@ def get_cookie_and_exit() -> None:
             print("\n🍪 To get your LinkedIn cookie manually:")
             print("   1. Login to LinkedIn in your browser")
             print("   2. Open Developer Tools (F12)")
-            print("   3. Go to Application/Storage > Cookies > linkedin.com")
+            print("   3. Go to Application/Storage > Cookies > www.linkedin.com")
             print("   4. Copy the 'li_at' cookie value")
             print("   5. Set LINKEDIN_COOKIE environment variable or use --cookie flag")
         elif "invalid credentials" in error_msg:
@@ -144,8 +148,9 @@ def ensure_authentication_ready() -> str:
     # If in non-interactive mode and no auth, fail immediately
     if config.chrome.non_interactive:
         raise CredentialsNotFoundError(
-            "No LinkedIn authentication found. Please provide cookie via "
-            "environment variable (LINKEDIN_COOKIE) or run with --get-cookie to obtain one."
+            "No LinkedIn cookie found for non-interactive mode. You can:\n"
+            "  1. Set LINKEDIN_COOKIE environment variable with a valid LinkedIn session cookie\n"
+            "  2. Run with --get-cookie to extract a cookie using email/password"
         )
 
     # Run interactive setup
@@ -211,9 +216,12 @@ def main() -> None:
         logger.info("Authentication ready")
     except CredentialsNotFoundError as e:
         logger.error(f"Authentication setup failed: {e}")
-        print(
-            "\n❌ Authentication required - please provide LinkedIn cookie or credentials"
-        )
+        if config.chrome.non_interactive:
+            print("\n❌ LinkedIn cookie required for Docker/non-interactive mode")
+        else:
+            print(
+                "\n❌ Authentication required - please provide LinkedIn authentication"
+            )
         sys.exit(1)
     except KeyboardInterrupt:
         print("\n\n👋 Setup cancelled by user")