diff --git a/cognee/api/v1/ui/ui.py b/cognee/api/v1/ui/ui.py index 6ad7b1cb4..a766529e9 100644 --- a/cognee/api/v1/ui/ui.py +++ b/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) diff --git a/cognee/cli/_cognee.py b/cognee/cli/_cognee.py index 0bef760aa..4fec89192 100644 --- a/cognee/cli/_cognee.py +++ b/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