From 7220052ca64c5ab28762d1144428db551d93d320 Mon Sep 17 00:00:00 2001
From: Daulet Amirkhanov <damirkhanov01@gmail.com>
Date: Thu, 11 Sep 2025 15:21:31 +0100
Subject: [PATCH] handle cli started ui closure gracefully

---
 cognee/api/v1/ui/ui.py | 35 ++++++++++++++++++++++++++++++++---
 cognee/cli/_cognee.py  | 41 ++++++++++++++++++++++++++++++++++-------
 2 files changed, 66 insertions(+), 10 deletions(-)

diff --git a/cognee/api/v1/ui/ui.py b/cognee/api/v1/ui/ui.py
index db876978e..6806e97ce 100644
--- a/cognee/api/v1/ui/ui.py
+++ b/cognee/api/v1/ui/ui.py
@@ -1,4 +1,5 @@
 import os
+import signal
 import subprocess
 import threading
 import time
@@ -401,6 +402,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
         process = subprocess.Popen(
             ["npm", "run", "dev"],
             cwd=frontend_path,
@@ -408,6 +410,7 @@ 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
         )
 
         # Give it a moment to start up
@@ -447,7 +450,7 @@ def start_ui(
 
 def stop_ui(process: subprocess.Popen) -> bool:
     """
-    Stop a running UI server process.
+    Stop a running UI server process and all its children.
 
     Args:
         process: The subprocess.Popen object returned by start_ui()
@@ -459,12 +462,38 @@ def stop_ui(process: subprocess.Popen) -> bool:
         return False
 
     try:
-        process.terminate()
+        # Try to terminate the process group (includes child processes like Next.js)
+        if hasattr(os, 'killpg'):
+            try:
+                # Kill the entire process group
+                os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+                logger.debug("Sent SIGTERM to process group")
+            except (OSError, ProcessLookupError):
+                # Fall back to terminating just the main process
+                process.terminate()
+                logger.debug("Terminated main process only")
+        else:
+            process.terminate()
+            logger.debug("Terminated main process (Windows)")
+        
         try:
             process.wait(timeout=10)
+            logger.info("UI server stopped gracefully")
         except subprocess.TimeoutExpired:
             logger.warning("Process didn't terminate gracefully, forcing kill")
-            process.kill()
+            
+            # Force kill the process group
+            if hasattr(os, 'killpg'):
+                try:
+                    os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+                    logger.debug("Sent SIGKILL to process group")
+                except (OSError, ProcessLookupError):
+                    process.kill()
+                    logger.debug("Force killed main process only")
+            else:
+                process.kill()
+                logger.debug("Force killed main process (Windows)")
+            
             process.wait()
 
         logger.info("UI server stopped")
diff --git a/cognee/cli/_cognee.py b/cognee/cli/_cognee.py
index d889cc78d..dab35aeb5 100644
--- a/cognee/cli/_cognee.py
+++ b/cognee/cli/_cognee.py
@@ -1,6 +1,8 @@
 import sys
 import os
 import argparse
+import signal
+import subprocess
 from typing import Any, Sequence, Dict, Type, cast, List
 import click
 
@@ -172,16 +174,43 @@ def main() -> int:
 
     # Handle UI flag
     if hasattr(args, 'start_ui') and args.start_ui:
+        server_process = None
+        
+        def signal_handler(signum, frame):
+            """Handle Ctrl+C and other termination signals"""
+            nonlocal server_process
+            fmt.echo("\nShutting down UI server...")
+            if server_process:
+                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}")
+            sys.exit(0)
+        
+        # Set up signal handlers
+        signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
+        signal.signal(signal.SIGTERM, signal_handler)  # Termination request
+        
         try:
             from cognee import start_ui
             fmt.echo("Starting cognee UI...")
-            server = start_ui(
+            server_process = start_ui(
                 host="localhost",
                 port=3001,
                 open_browser=True
             )
             
-            if server:
+            if server_process:
                 fmt.success("UI server started successfully!")
                 fmt.echo("The interface is available at: http://localhost:3001")
                 fmt.note("Press Ctrl+C to stop the server...")
@@ -189,13 +218,11 @@ def main() -> int:
                 try:
                     # Keep the server running
                     import time
-                    while server.poll() is None:  # While process is still running
+                    while server_process.poll() is None:  # While process is still running
                         time.sleep(1)
                 except KeyboardInterrupt:
-                    fmt.echo("\nStopping UI server...")
-                    server.terminate()
-                    server.wait()
-                    fmt.success("UI server stopped.")
+                    # This shouldn't happen now due to signal handler, but kept for safety
+                    signal_handler(signal.SIGINT, None)
                 
                 return 0
             else: