handle cli started ui closure gracefully

This commit is contained in:
Daulet Amirkhanov 2025-09-11 15:21:31 +01:00
parent d93f31ad30
commit 7220052ca6
2 changed files with 66 additions and 10 deletions

View file

@ -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")

View file

@ -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: