fix(runtime): validate engine installation completion and preserve live install state (fixes #112)

backend/app/core/task_registry.py

@@ -22,6 +22,7 @@ def __init__(self) -> None:
         self._tasks: Dict[str, Dict[str, Any]] = {}
         self._cancellations: Dict[str, asyncio.Event] = {}
         self._update_locks: Dict[str, asyncio.Lock] = {}
+        self._install_locks: Dict[str, asyncio.Lock] = {}
 
     def register_task(
         self,
@@ -83,5 +84,14 @@ def get_update_lock(self, engine_key: str) -> asyncio.Lock:
             self._update_locks[engine_key] = asyncio.Lock()
         return self._update_locks[engine_key]
 
+    def get_install_lock(self, engine_key: str) -> asyncio.Lock:
+        """Retrieve or create an exclusive async lease lock for an engine installation."""
+        if not hasattr(self, "_install_locks"):
+            self._install_locks: Dict[str, asyncio.Lock] = {}
+        if engine_key not in self._install_locks:
+            self._install_locks[engine_key] = asyncio.Lock()
+        return self._install_locks[engine_key]
+
 
 task_registry = TaskLifecycleRegistry()
+

backend/app/runtime/installer.py

@@ -37,6 +37,9 @@
 class IsolatedEngineInstaller:
     """Manages staged installation of local engines inside hermetic environments."""
 
+    # Class-level set of currently active installation tasks (engine_type.value)
+    _active_installs: set[str] = set()
+
     def __init__(self, engine_dir: Optional[Path] = None) -> None:
         self.engine_dir = engine_dir or get_default_engine_dir()
 
@@ -55,14 +58,17 @@ def read_manifest(self, engine_type: EngineType) -> EngineInstallManifest:
                 data = json.loads(manifest_path.read_text(encoding="utf-8"))
                 manifest = EngineInstallManifest.model_validate(data)
                 # Check for uncompleted previous run (interruption detection)
+                # Only mark as INTERRUPTED if there is NO active in-memory install worker running!
                 if manifest.phase in (
+                    InstallPhase.CHECKING,
                     InstallPhase.CREATING_VENV,
                     InstallPhase.DOWNLOADING,
                     InstallPhase.INSTALLING_DEPS,
                 ):
-                    manifest.phase = InstallPhase.INTERRUPTED
-                    manifest.error_message = "Installation was interrupted before completion."
-                    self.write_manifest(manifest)
+                    if engine_type.value not in self._active_installs:
+                        manifest.phase = InstallPhase.INTERRUPTED
+                        manifest.error_message = "Installation was interrupted before completion."
+                        self.write_manifest(manifest)
                 return manifest
             except Exception as e:
                 logger.warning(f"Error reading manifest for {engine_type}: {e}")
@@ -91,6 +97,34 @@ def _get_pip_bin(self, runtime_dir: Path) -> Path:
             return runtime_dir / "Scripts" / "pip.exe"
         return runtime_dir / "bin" / "pip"
 
+    def _get_expected_entrypoints(self, engine_type: EngineType, engine_target: Path) -> list[Path]:
+        """Return expected entrypoint script paths for the engine."""
+        if engine_type == EngineType.COMFYUI:
+            return [engine_target / "main.py"]
+        elif engine_type == EngineType.WEBUI:
+            return [engine_target / "launch.py", engine_target / "webui.py"]
+        return []
+
+    def validate_installation(self, engine_type: EngineType, engine_target: Path, runtime_target: Path) -> None:
+        """
+        Validate that the engine installation is complete and healthy:
+        1. Virtual environment Python binary exists
+        2. Expected entrypoint script exists
+        3. Engine requirements.txt exists
+        """
+        python_bin = self._get_python_bin(runtime_target)
+        if not python_bin.is_file():
+            raise RuntimeError(f"Hermetic Python binary missing at {python_bin}")
+
+        entrypoints = self._get_expected_entrypoints(engine_type, engine_target)
+        if not any(ep.is_file() for ep in entrypoints):
+            expected_names = " or ".join(ep.name for ep in entrypoints) or "entrypoint"
+            raise RuntimeError(f"Engine entrypoint ({expected_names}) missing in {engine_target}")
+
+        req_file = engine_target / "requirements.txt"
+        if not req_file.is_file():
+            raise RuntimeError(f"Engine requirements.txt missing in {engine_target}")
+
     async def create_isolated_venv(
         self, runtime_dir: Path, on_log: Optional[Callable[[str], None]] = None
     ) -> bool:
@@ -120,68 +154,94 @@ async def install_engine(
     ) -> EngineInstallManifest:
         """
         Run complete isolated installation:
-        1. Check environment & disk space
-        2. Create isolated sandboxed venv
-        3. Clone / stage repository code
+        1. Acquire per-engine install lock
+        2. Check environment & create isolated sandboxed venv
+        3. Clone / stage repository code into staging or target
         4. Install isolated dependencies via sandboxed pip
+        5. Validate entrypoint and health before marking completed
         """
-        manifest = self.read_manifest(engine_type)
-        now_str = datetime.now(timezone.utc).isoformat()
-        manifest.created_at = now_str
-        manifest.phase = InstallPhase.CHECKING
-        self.write_manifest(manifest)
-        if on_progress:
-            on_progress(manifest)
-
-        engine_target = Path(manifest.engine_dir)
-        runtime_target = Path(manifest.runtime_dir)
-
-        try:
-            # Phase 1: Virtualenv Creation
-            manifest.phase = InstallPhase.CREATING_VENV
-            manifest.last_log_line = "Initializing sandboxed virtual environment..."
-            self.write_manifest(manifest)
-            if on_progress:
-                on_progress(manifest)
-
-            venv_ok = await self.create_isolated_venv(runtime_target)
-            if not venv_ok:
-                raise RuntimeError(f"Failed to create virtual environment at {runtime_target}")
-
-            python_bin = self._get_python_bin(runtime_target)
-            manifest.python_bin = str(python_bin)
-
-            # Phase 2: Stage / Download code
-            manifest.phase = InstallPhase.DOWNLOADING
-            manifest.last_log_line = f"Staging {engine_type.value} repository..."
-            self.write_manifest(manifest)
-            if on_progress:
-                on_progress(manifest)
-
-            repo_url = COMFYUI_GIT_REPO if engine_type == EngineType.COMFYUI else WEBUI_GIT_REPO
-            if not engine_target.exists():
-                clone_cmd = ["git", "clone", "--depth", "1", repo_url, str(engine_target)]
-                proc = await asyncio.create_subprocess_exec(
-                    *clone_cmd,
-                    stdout=asyncio.subprocess.PIPE,
-                    stderr=asyncio.subprocess.PIPE,
-                )
-                stdout, stderr = await proc.communicate()
-                if proc.returncode != 0:
-                    err_msg = stderr.decode(errors="replace").strip()
-                    raise RuntimeError(f"Git clone failed: {err_msg}")
+        engine_key = engine_type.value
+        install_lock = task_registry.get_install_lock(engine_key)
 
-            # Phase 3: Install isolated dependencies
-            manifest.phase = InstallPhase.INSTALLING_DEPS
-            manifest.last_log_line = "Installing engine requirements into isolated environment..."
+        if install_lock.locked():
+            manifest = self.read_manifest(engine_type)
+            manifest.phase = InstallPhase.FAILED
+            manifest.error_message = f"An installation lease is already held for {engine_type.value}."
             self.write_manifest(manifest)
-            if on_progress:
-                on_progress(manifest)
+            return manifest
 
-            req_file = engine_target / "requirements.txt"
-            pip_bin = self._get_pip_bin(runtime_target)
+        async with install_lock:
+            self._active_installs.add(engine_key)
+            try:
+                manifest = self.read_manifest(engine_type)
+                now_str = datetime.now(timezone.utc).isoformat()
+                manifest.created_at = now_str
+                manifest.phase = InstallPhase.CHECKING
+                self.write_manifest(manifest)
+                if on_progress:
+                    on_progress(manifest)
+
+                engine_target = Path(manifest.engine_dir)
+                runtime_target = Path(manifest.runtime_dir)
+
+                # Phase 1: Virtualenv Creation
+                manifest.phase = InstallPhase.CREATING_VENV
+                manifest.last_log_line = "Initializing sandboxed virtual environment..."
+                self.write_manifest(manifest)
+                if on_progress:
+                    on_progress(manifest)
+
+                venv_ok = await self.create_isolated_venv(runtime_target)
+                if not venv_ok:
+                    raise RuntimeError(f"Failed to create virtual environment at {runtime_target}")
+
+                python_bin = self._get_python_bin(runtime_target)
+                manifest.python_bin = str(python_bin)
+
+                # Phase 2: Stage / Download code
+                manifest.phase = InstallPhase.DOWNLOADING
+                manifest.last_log_line = f"Staging {engine_type.value} repository..."
+                self.write_manifest(manifest)
+                if on_progress:
+                    on_progress(manifest)
+
+                repo_url = COMFYUI_GIT_REPO if engine_type == EngineType.COMFYUI else WEBUI_GIT_REPO
+                entrypoints = self._get_expected_entrypoints(engine_type, engine_target)
+
+                # If destination directory exists but has no valid entrypoint or git repo, clean it up before cloning
+                if engine_target.exists():
+                    has_entrypoint = any(ep.is_file() for ep in entrypoints)
+                    has_git = (engine_target / ".git").is_dir()
+                    if not has_entrypoint or not has_git:
+                        shutil.rmtree(engine_target, ignore_errors=True)
+
+                if not engine_target.exists():
+                    clone_cmd = ["git", "clone", "--depth", "1", repo_url, str(engine_target)]
+                    proc = await asyncio.create_subprocess_exec(
+                        *clone_cmd,
+                        stdout=asyncio.subprocess.PIPE,
+                        stderr=asyncio.subprocess.PIPE,
+                    )
+                    stdout, stderr = await proc.communicate()
+                    if proc.returncode != 0:
+                        err_msg = stderr.decode(errors="replace").strip()
+                        raise RuntimeError(f"Git clone failed: {err_msg}")
+
+                # Phase 3: Install isolated dependencies
+                manifest.phase = InstallPhase.INSTALLING_DEPS
+                manifest.last_log_line = "Installing engine requirements into isolated environment..."
+                self.write_manifest(manifest)
+                if on_progress:
+                    on_progress(manifest)
+
+                req_file = engine_target / "requirements.txt"
+                pip_bin = self._get_pip_bin(runtime_target)
+
+                if not req_file.is_file():
+                    raise RuntimeError(f"requirements.txt missing in {engine_target}")
+                if not pip_bin.is_file():
+                    raise RuntimeError(f"Hermetic pip binary missing in {runtime_target}")
 
-            if req_file.is_file() and pip_bin.is_file():
                 # Strict: execute ONLY pip inside runtime_target!
                 install_cmd = [str(pip_bin), "install", "--no-warn-script-location", "-r", str(req_file)]
                 proc = await asyncio.create_subprocess_exec(
@@ -195,22 +255,27 @@ async def install_engine(
                     err_msg = stderr.decode(errors="replace").strip()
                     raise RuntimeError(f"Pip install failed: {err_msg[:300]}")
 
-            # Completed!
-            manifest.phase = InstallPhase.COMPLETED
-            manifest.completed_at = datetime.now(timezone.utc).isoformat()
-            manifest.last_log_line = f"{engine_type.value} installed successfully."
-            manifest.error_message = None
-            self.write_manifest(manifest)
-            if on_progress:
-                on_progress(manifest)
+                # Phase 4: Validation
+                self.validate_installation(engine_type, engine_target, runtime_target)
 
-        except Exception as e:
-            manifest.phase = InstallPhase.FAILED
-            manifest.error_message = str(e)
-            manifest.last_log_line = f"Installation error: {e}"
-            self.write_manifest(manifest)
-            if on_progress:
-                on_progress(manifest)
+                # Completed!
+                manifest.phase = InstallPhase.COMPLETED
+                manifest.completed_at = datetime.now(timezone.utc).isoformat()
+                manifest.last_log_line = f"{engine_type.value} installed successfully."
+                manifest.error_message = None
+                self.write_manifest(manifest)
+                if on_progress:
+                    on_progress(manifest)
+
+            except Exception as e:
+                manifest.phase = InstallPhase.FAILED
+                manifest.error_message = str(e)
+                manifest.last_log_line = f"Installation error: {e}"
+                self.write_manifest(manifest)
+                if on_progress:
+                    on_progress(manifest)
+            finally:
+                self._active_installs.discard(engine_key)
 
         return manifest
 

backend/tests/test_m2_engines.py

@@ -20,6 +20,7 @@
 from fastapi.testclient import TestClient
 
 from app.main import app
+from app.core.task_registry import task_registry
 from app.runtime.engine_manager import EngineManager
 from app.runtime.hardware import check_hardware_readiness, detect_gpus, get_storage_readiness
 from app.runtime.installer import IsolatedEngineInstaller
@@ -90,7 +91,7 @@ def test_installer_manifest_and_interruption():
         assert manifest.phase == InstallPhase.IDLE
         assert manifest.engine_type == EngineType.COMFYUI
 
-        # 2. Simulate an interrupted install state
+        # 2. Simulate an interrupted install state (no active in-memory task)
         manifest.phase = InstallPhase.DOWNLOADING
         installer_inst.write_manifest(manifest)
 
@@ -99,6 +100,59 @@ def test_installer_manifest_and_interruption():
         assert reloaded.phase == InstallPhase.INTERRUPTED
         assert "interrupted" in (reloaded.error_message or "").lower()
 
+        # 4. Live polling while an install is actively running must NOT mark as interrupted
+        installer_inst._active_installs.add(EngineType.COMFYUI.value)
+        try:
+            manifest.phase = InstallPhase.DOWNLOADING
+            manifest.error_message = None
+            installer_inst.write_manifest(manifest)
+            live_manifest = installer_inst.read_manifest(EngineType.COMFYUI)
+            assert live_manifest.phase == InstallPhase.DOWNLOADING
+            assert live_manifest.error_message is None
+        finally:
+            installer_inst._active_installs.discard(EngineType.COMFYUI.value)
+
+
+@pytest.mark.asyncio
+async def test_installer_empty_dir_missing_entrypoint_rejection(tmp_path):
+    """Verify that an empty engine directory or missing entrypoint is rejected and does not report COMPLETED."""
+    installer_inst = IsolatedEngineInstaller(engine_dir=tmp_path)
+    engine_target = tmp_path / "comfyui"
+    engine_target.mkdir(parents=True, exist_ok=True)
+    runtime_target = tmp_path / "runtime"
+
+    # Mock venv creation as succeeding and creating python binary
+    python_bin = installer_inst._get_python_bin(runtime_target)
+    python_bin.parent.mkdir(parents=True, exist_ok=True)
+    python_bin.write_text("# mock python", encoding="utf-8")
+    pip_bin = installer_inst._get_pip_bin(runtime_target)
+    pip_bin.write_text("# mock pip", encoding="utf-8")
+
+    with patch.object(installer_inst, "create_isolated_venv", return_value=True):
+        # Empty directory with no main.py and no .git should trigger cleanup & clone attempt,
+        # but mock git clone to do nothing, simulating missing entrypoint failure
+        with patch("asyncio.create_subprocess_exec") as mock_exec:
+            proc_mock = AsyncMock()
+            proc_mock.returncode = 0
+            proc_mock.communicate.return_value = (b"", b"")
+            mock_exec.return_value = proc_mock
+
+            manifest = await installer_inst.install_engine(EngineType.COMFYUI)
+            assert manifest.phase == InstallPhase.FAILED
+            assert "requirements.txt missing" in (manifest.error_message or "") or "entrypoint" in (manifest.error_message or "")
+
+
+@pytest.mark.asyncio
+async def test_installer_concurrent_lease_lock_protection(tmp_path):
+    """Verify that concurrent install calls for the same engine are rejected by lease lock."""
+    installer_inst = IsolatedEngineInstaller(engine_dir=tmp_path)
+    lock = task_registry.get_install_lock(EngineType.COMFYUI.value)
+
+    async with lock:
+        manifest = await installer_inst.install_engine(EngineType.COMFYUI)
+        assert manifest.phase == InstallPhase.FAILED
+        assert "lease is already held" in (manifest.error_message or "")
+
 
 @pytest.mark.asyncio
 async def test_external_engine_connection_zero_ownership():