diff --git a/cognee/api/v1/ui/ui.py b/cognee/api/v1/ui/ui.py index 2c3687e16..d0f109393 100644 --- a/cognee/api/v1/ui/ui.py +++ b/cognee/api/v1/ui/ui.py @@ -502,20 +502,22 @@ def start_ui( if start_mcp: logger.info("Starting Cognee MCP server with Docker...") - cwd = os.getcwd() - env_file = os.path.join(cwd, ".env") try: image = "cognee/cognee-mcp:main" subprocess.run(["docker", "pull", image], check=True) + import uuid + + container_name = f"cognee-mcp-{uuid.uuid4().hex[:8]}" + docker_cmd = [ "docker", "run", + "--name", + container_name, "-p", f"{mcp_port}:8000", "--rm", - "--env-file", - env_file, "-e", "TRANSPORT_MODE=sse", ] @@ -531,6 +533,10 @@ def start_ui( f"Configuring MCP to connect to backend API at http://localhost:{backend_port}" ) logger.info("(localhost will be auto-converted to host.docker.internal)") + else: + cwd = os.getcwd() + env_file = os.path.join(cwd, ".env") + docker_cmd.extend(["--env-file", env_file]) docker_cmd.append("cognee/cognee-mcp:main") @@ -544,7 +550,8 @@ def start_ui( _stream_process_output(mcp_process, "stdout", "[MCP]", "\033[34m") # Blue _stream_process_output(mcp_process, "stderr", "[MCP]", "\033[34m") # Blue - pid_callback(mcp_process.pid) + # Pass both PID and container name using a tuple + pid_callback((mcp_process.pid, container_name)) mode_info = "API mode" if start_backend else "direct mode" logger.info( diff --git a/cognee/cli/_cognee.py b/cognee/cli/_cognee.py index b257d37de..9008d68e3 100644 --- a/cognee/cli/_cognee.py +++ b/cognee/cli/_cognee.py @@ -175,12 +175,40 @@ def main() -> int: # Handle UI flag if hasattr(args, "start_ui") and args.start_ui: spawned_pids = [] + docker_container = None def signal_handler(signum, frame): """Handle Ctrl+C and other termination signals""" - nonlocal spawned_pids + nonlocal spawned_pids, docker_container fmt.echo("\nShutting down UI server...") + # First, stop Docker container if running + if docker_container: + try: + fmt.echo(f"Stopping Docker container {docker_container}...") + result = subprocess.run( + ["docker", "stop", docker_container], + capture_output=True, + timeout=10, + check=False, + ) + if result.returncode == 0: + fmt.success(f"✓ Docker container {docker_container} stopped.") + else: + fmt.warning( + f"Could not stop container {docker_container}: {result.stderr.decode()}" + ) + except subprocess.TimeoutExpired: + fmt.warning( + f"Timeout stopping container {docker_container}, forcing removal..." + ) + subprocess.run( + ["docker", "rm", "-f", docker_container], capture_output=True, check=False + ) + except Exception as e: + fmt.warning(f"Could not stop Docker container {docker_container}: {e}") + + # Then, stop regular processes for pid in spawned_pids: try: if hasattr(os, "killpg"): @@ -209,10 +237,16 @@ 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) + # Callback to capture PIDs and Docker container of all spawned processes + def pid_callback(pid_or_tuple): + nonlocal spawned_pids, docker_container + # Handle both regular PIDs and (PID, container_name) tuples + if isinstance(pid_or_tuple, tuple): + pid, container_name = pid_or_tuple + spawned_pids.append(pid) + docker_container = container_name + else: + spawned_pids.append(pid_or_tuple) frontend_port = 3000 start_backend, backend_port = True, 8000