add more robust process cleanup to avoid dangling subprocesses

cognee/api/v1/ui/ui.py

@@ -7,7 +7,7 @@ import webbrowser
 import zipfile
 import requests
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Callable, Optional, Tuple
 import tempfile
 import shutil
 
@@ -326,6 +326,7 @@ def prompt_user_for_download() -> bool:
 
 
 def start_ui(
+    pid_callback: Callable[[int], None],
     host: str = "localhost",
     port: int = 3000,
     open_browser: bool = True,
@@ -346,6 +347,7 @@ def start_ui(
     6. Optionally open the browser
 
     Args:
+        pid_callback: Callback to notify with PID of each spawned process
         host: Host to bind the frontend server to (default: localhost)
         port: Port to run the frontend server on (default: 3000)
         open_browser: Whether to open the browser automatically (default: True)
@@ -397,6 +399,8 @@ def start_ui(
                 preexec_fn=os.setsid if hasattr(os, "setsid") else None,
             )
 
+            pid_callback(backend_process.pid)
+
             # Give the backend a moment to start
             time.sleep(2)
 
@@ -460,7 +464,7 @@ def start_ui(
     logger.info("This may take a moment to compile and start...")
 
     try:
-        # Use process group to ensure all child processes get terminated together
+        # Create frontend in its own process group for clean termination
         process = subprocess.Popen(
             ["npm", "run", "dev"],
             cwd=frontend_path,
@@ -468,11 +472,11 @@ def start_ui(
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             text=True,
-            preexec_fn=os.setsid
-            if hasattr(os, "setsid")
-            else None,  # Create new process group on Unix
+            preexec_fn=os.setsid if hasattr(os, "setsid") else None,
         )
 
+        pid_callback(process.pid)
+
         # Give it a moment to start up
         time.sleep(3)
 

cognee/cli/_cognee.py

@@ -174,30 +174,23 @@ def main() -> int:
 
     # Handle UI flag
     if hasattr(args, "start_ui") and args.start_ui:
-        server_process = None
+        spawned_pids = []
 
         def signal_handler(signum, frame):
             """Handle Ctrl+C and other termination signals"""
-            nonlocal server_process
+            nonlocal spawned_pids
             fmt.echo("\nShutting down UI server...")
-            if server_process:
+
+            for pid in spawned_pids:
                 try:
-                    # Try graceful termination first
-                    server_process.terminate()
-                    try:
-                        server_process.wait(timeout=5)
-                        fmt.success("UI server stopped gracefully.")
-                    except subprocess.TimeoutExpired:
-                        # If graceful termination fails, force kill
-                        fmt.echo("Force stopping UI server...")
-                        server_process.kill()
-                        server_process.wait()
-                        fmt.success("UI server stopped.")
-                except Exception as e:
-                    fmt.warning(f"Error stopping server: {e}")
+                    pgid = os.getpgid(pid)
+                    os.killpg(pgid, signal.SIGTERM)
+                    fmt.success(f"✓ Process group {pgid} (PID {pid}) terminated.")
+                except (OSError, ProcessLookupError) as e:
+                    fmt.warning(f"Could not terminate process {pid}: {e}")
+
             sys.exit(0)
 
-        # Set up signal handlers
         signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
         signal.signal(signal.SIGTERM, signal_handler)  # Termination request
 
@@ -206,8 +199,17 @@ def main() -> int:
 
             fmt.echo("Starting cognee UI...")
 
+            # Callback to capture PIDs of all spawned processes
+            def pid_callback(pid):
+                nonlocal spawned_pids
+                spawned_pids.append(pid)
+
             server_process = start_ui(
-                host="localhost", port=3000, open_browser=True, start_backend=True
+                host="localhost",
+                port=3000,
+                open_browser=True,
+                start_backend=True,
+                pid_callback=pid_callback,
             )
 
             if server_process:
@@ -229,10 +231,12 @@ def main() -> int:
                 return 0
             else:
                 fmt.error("Failed to start UI server. Check the logs above for details.")
+                signal_handler(signal.SIGTERM, None)
                 return 1
 
         except Exception as ex:
             fmt.error(f"Error starting UI: {str(ex)}")
+            signal_handler(signal.SIGTERM, None)
             if debug.is_debug_enabled():
                 raise ex
             return 1