refactor(drivers): improve session management and cleanup

Author: stickerdaniel

linkedin_mcp_server/drivers/chrome.py

@@ -8,6 +8,8 @@
 
 import logging
 import os
+import shutil
+import tempfile
 from typing import Dict, Optional
 
 from linkedin_scraper.exceptions import (
@@ -32,13 +34,19 @@
 # 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() -> webdriver.Chrome:
+def create_chrome_driver(session_id: str = "default") -> webdriver.Chrome:
     """
     Create a new Chrome WebDriver instance with proper configuration.
 
+    Args:
+        session_id: Unique identifier for the session (used for cleanup)
+
     Returns:
         webdriver.Chrome: Configured Chrome WebDriver instance
 
@@ -63,6 +71,14 @@ def create_chrome_driver() -> webdriver.Chrome:
     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
+
     # Set user agent (configurable with sensible default)
     user_agent = getattr(config.chrome, "user_agent", DEFAULT_USER_AGENT)
     chrome_options.add_argument(f"--user-agent={user_agent}")
@@ -213,7 +229,7 @@ def get_or_create_driver(authentication: str) -> webdriver.Chrome:
 
     try:
         # Create new driver
-        driver = create_chrome_driver()
+        driver = create_chrome_driver(session_id)
 
         # Login to LinkedIn
         login_to_linkedin(driver, authentication)
@@ -245,7 +261,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
+    global active_drivers, user_data_dirs
 
     for session_id, driver in active_drivers.items():
         try:
@@ -254,8 +270,21 @@ 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()
-    logger.info("All Chrome WebDriver sessions closed")
+    user_data_dirs.clear()
+    logger.info("All Chrome WebDriver sessions closed and cleaned up")
 
 
 def get_active_driver() -> Optional[webdriver.Chrome]:

linkedin_mcp_server/server.py

@@ -10,7 +10,6 @@
 
 from fastmcp import FastMCP
 
-from linkedin_mcp_server.drivers.chrome import active_drivers
 from linkedin_mcp_server.tools.company import register_company_tools
 from linkedin_mcp_server.tools.job import register_job_tools
 from linkedin_mcp_server.tools.person import register_person_tools
@@ -31,42 +30,25 @@ def create_mcp_server() -> FastMCP:
     @mcp.tool()
     async def close_session() -> Dict[str, Any]:
         """Close the current browser session and clean up resources."""
-        session_id = "default"  # Using the same default session
+        from linkedin_mcp_server.drivers.chrome import close_all_drivers
 
-        if session_id in active_drivers:
-            try:
-                active_drivers[session_id].quit()
-                del active_drivers[session_id]
-                return {
-                    "status": "success",
-                    "message": "Successfully closed the browser session",
-                }
-            except Exception as e:
-                return {
-                    "status": "error",
-                    "message": f"Error closing browser session: {str(e)}",
-                }
-        else:
+        try:
+            close_all_drivers()
+            return {
+                "status": "success",
+                "message": "Successfully closed the browser session and cleaned up resources",
+            }
+        except Exception as e:
             return {
-                "status": "warning",
-                "message": "No active browser session to close",
+                "status": "error",
+                "message": f"Error closing browser session: {str(e)}",
             }
 
     return mcp
 
 
 def shutdown_handler() -> None:
     """Clean up resources on shutdown."""
-    for session_id, driver in list(active_drivers.items()):
-        try:
-            driver.quit()
-            del active_drivers[session_id]
-        except Exception as e:
-            logger.error(
-                f"Error closing driver during shutdown: {e}",
-                extra={
-                    "session_id": session_id,
-                    "exception_type": type(e).__name__,
-                    "exception_message": str(e),
-                },
-            )
+    from linkedin_mcp_server.drivers.chrome import close_all_drivers
+
+    close_all_drivers()

linkedin_mcp_server/setup.py

@@ -7,6 +7,7 @@
 
 import logging
 import os
+import tempfile
 from contextlib import contextmanager
 from typing import Dict, Iterator
 
@@ -119,6 +120,11 @@ def temporary_chrome_driver() -> Iterator[webdriver.Chrome]:
     # 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}")
+
     driver = None
     try:
         # Create temporary driver