fix: reliable, self-healing AI feature-bundle installs (#472)

Make on-demand AI feature-bundle installs reliable and self-healing, closing
the failure modes behind most "some tool doesn't work" reports.

Multi-bundle installs: tools needing more than one bundle (Passport Photo,
Enhance Faces) install every required bundle from one action and stay
not-installed until all are present. Verified across all 19 AI tools.

Downloads: self-heal the accelerated Hugging Face (Xet) client so an upgraded
venv no longer silently falls back to slow urllib; restart instead of
corrupting a resumed partial when a proxy ignores Range and returns 200;
verify the completed size; fail fast on disk-full and HTTP 4xx; retry
transient errors five times; add hf_transfer fallback and document Xet egress.

Install integrity: crash-atomic venv writes so a killed or out-of-space
install can no longer tear the shared venv and break other tools; a boot
breadcrumb reseeds a torn venv to a clean state automatically; a post-install
smoke import test refuses to record a bundle whose libraries cannot load; an
install watchdog stops a wedged installer that would otherwise hold the venv
writer lock forever.

Adds unit and end-to-end tests for every failure mode above.

packages/ai/python/tests/test_install_feature_download.py

@@ -0,0 +1,151 @@
+import importlib.util
+import os
+import sys
+import types
+
+
+def load_installer():
+    script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
+    spec = importlib.util.spec_from_file_location("install_feature_under_test", script_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_download_with_hf_hub_uses_accelerated_client(monkeypatch, tmp_path):
+    installer = load_installer()
+    downloaded = tmp_path / "hf-cache" / "bundle.tar.gz"
+    downloaded.parent.mkdir()
+
+    calls = {}
+
+    def fake_hf_hub_download(**kwargs):
+        calls.update(kwargs)
+        downloaded.write_bytes(b"archive")
+        return str(downloaded)
+
+    fake_module = types.ModuleType("huggingface_hub")
+    fake_module.hf_hub_download = fake_hf_hub_download
+    monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
+
+    progress = []
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: progress.append((p, s)))
+
+    dest = tmp_path / "staging" / "object-eraser-colorize-amd64-gpu.tar.gz"
+    dest.parent.mkdir()
+
+    assert (
+        installer.download_with_hf_hub(
+            "snapotter/feature-bundles",
+            "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
+            str(dest),
+            100,
+            2,
+            85,
+        )
+        is True
+    )
+
+    assert dest.read_bytes() == b"archive"
+    assert calls["repo_id"] == "snapotter/feature-bundles"
+    assert calls["repo_type"] == "model"
+    assert calls["filename"] == "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz"
+    assert any("accelerated" in stage.lower() for _, stage in progress)
+
+
+def test_download_with_hf_hub_cleans_cache_when_download_raises(monkeypatch, tmp_path):
+    """A failed accelerated download must not leak its .cache staging tree onto
+    disk before the urllib fallback runs."""
+    installer = load_installer()
+    staging = tmp_path / "staging"
+    staging.mkdir()
+    # Simulate a partial hf cache tree left behind by a failed transfer.
+    leaked_cache = staging / ".cache" / "huggingface" / "download"
+    leaked_cache.mkdir(parents=True)
+    (leaked_cache / "blob.incomplete").write_bytes(b"partial")
+    leaked_nested = staging / "v2.0.0"
+    leaked_nested.mkdir()
+
+    def fake_hf_hub_download(**_kwargs):
+        raise RuntimeError("xet CAS unreachable")
+
+    fake_module = types.ModuleType("huggingface_hub")
+    fake_module.hf_hub_download = fake_hf_hub_download
+    monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+
+    dest = staging / "object-eraser-colorize-amd64-gpu.tar.gz"
+
+    assert (
+        installer.download_with_hf_hub(
+            "deepsafe/feature-bundles",
+            "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
+            str(dest),
+            100,
+            2,
+            85,
+        )
+        is False
+    )
+    # Both the .cache tree and the nested archive dir are reclaimed.
+    assert not (staging / ".cache").exists()
+    assert not (staging / "v2.0.0").exists()
+
+
+def test_ensure_hf_hub_noops_when_client_already_importable(monkeypatch, tmp_path):
+    installer = load_installer()
+    fake_module = types.ModuleType("huggingface_hub")
+    monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
+
+    ran = {"pip": False}
+    monkeypatch.setattr(
+        installer.subprocess, "run", lambda *a, **k: ran.__setitem__("pip", True)
+    )
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+
+    installer.ensure_hf_hub(str(tmp_path))
+    assert ran["pip"] is False
+
+
+def test_ensure_hf_hub_self_heals_missing_client(monkeypatch, tmp_path):
+    """On a drifted venv where huggingface_hub is missing, ensure_hf_hub must
+    pip-install it into that venv rather than let the caller fall back to the
+    slow single-stream urllib downloader silently."""
+    installer = load_installer()
+    monkeypatch.delitem(sys.modules, "huggingface_hub", raising=False)
+
+    # Make the huggingface_hub import fail deterministically so ensure_hf_hub
+    # takes its self-heal branch.
+    import builtins
+
+    real_import = builtins.__import__
+
+    def blocked_import(name, *a, **k):
+        if name == "huggingface_hub":
+            raise ImportError("No module named 'huggingface_hub'")
+        return real_import(name, *a, **k)
+
+    monkeypatch.setattr(builtins, "__import__", blocked_import)
+
+    venv = tmp_path / "venv"
+    (venv / "bin").mkdir(parents=True)
+    (venv / "bin" / "python3").write_text("")
+
+    pip_calls = []
+    monkeypatch.setattr(
+        installer.subprocess,
+        "run",
+        lambda cmd, **k: pip_calls.append(cmd) or types.SimpleNamespace(returncode=0),
+    )
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+
+    installer.ensure_hf_hub(str(venv))
+
+    assert len(pip_calls) == 1
+    cmd = pip_calls[0]
+    assert cmd[0] == str(venv / "bin" / "python3")
+    assert "install" in cmd
+    spec = next(part for part in cmd if part.startswith("huggingface-hub["))
+    assert "hf_xet" in spec
+    assert "hf_transfer" in spec

packages/ai/python/tests/test_install_feature_move.py

@@ -0,0 +1,134 @@
+"""Crash-atomicity tests for move_tree, which writes into the SHARED venv.
+
+The invariant under test: a crash (or ENOSPC) partway through move_tree must
+never leave a destination entry missing. Because the venv is shared by every AI
+tool, a half-replaced package is what tears the venv and breaks unrelated tools.
+"""
+
+import errno
+import importlib.util
+import os
+
+import pytest
+
+
+def load_installer():
+    script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
+    spec = importlib.util.spec_from_file_location("install_feature_move_under_test", script_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_file_over_file_is_replaced_without_a_delete_window(monkeypatch, tmp_path):
+    installer = load_installer()
+    src = tmp_path / "src"
+    dst = tmp_path / "dst"
+    src.mkdir()
+    dst.mkdir()
+    (dst / "pkg.py").write_text("OLD")
+    (src / "pkg.py").write_text("NEW")
+
+    removed = []
+    real_remove = os.remove
+    monkeypatch.setattr(installer.os, "remove", lambda p: removed.append(p) or real_remove(p))
+
+    installer.move_tree(str(src), str(dst))
+
+    assert (dst / "pkg.py").read_text() == "NEW"
+    # The old file was atomically replaced, never deleted-then-rewritten.
+    assert str(dst / "pkg.py") not in removed
+
+
+def test_crash_mid_move_leaves_every_dest_old_or_new_never_missing(monkeypatch, tmp_path):
+    installer = load_installer()
+    src = tmp_path / "src"
+    dst = tmp_path / "dst"
+    src.mkdir()
+    dst.mkdir()
+    for name in ("a.py", "b.py", "c.py"):
+        (dst / name).write_text(f"OLD_{name}")
+        (src / name).write_text(f"NEW_{name}")
+
+    real_replace = os.replace
+    state = {"n": 0}
+
+    def flaky_replace(s, d):
+        state["n"] += 1
+        if state["n"] == 2:
+            raise OSError("simulated crash mid-move")
+        return real_replace(s, d)
+
+    monkeypatch.setattr(installer.os, "replace", flaky_replace)
+
+    with pytest.raises(OSError):
+        installer.move_tree(str(src), str(dst))
+
+    # Regardless of listdir order, every destination file must still exist and
+    # hold either its old or its new content -- never a torn/missing entry.
+    for name in ("a.py", "b.py", "c.py"):
+        assert (dst / name).exists()
+        assert (dst / name).read_text() in (f"OLD_{name}", f"NEW_{name}")
+
+
+def test_merges_directories_and_overwrites_files(tmp_path):
+    installer = load_installer()
+    src = tmp_path / "src"
+    dst = tmp_path / "dst"
+    (src / "pkg").mkdir(parents=True)
+    (dst / "pkg").mkdir(parents=True)
+    (dst / "pkg" / "keep.py").write_text("KEEP")
+    (src / "pkg" / "keep.py").write_text("UPDATED")
+    (src / "pkg" / "new.py").write_text("NEW")
+
+    installer.move_tree(str(src), str(dst))
+
+    assert (dst / "pkg" / "keep.py").read_text() == "UPDATED"
+    assert (dst / "pkg" / "new.py").read_text() == "NEW"
+
+
+def test_type_mismatch_dir_replaces_file(tmp_path):
+    installer = load_installer()
+    src = tmp_path / "src"
+    dst = tmp_path / "dst"
+    src.mkdir()
+    dst.mkdir()
+    (dst / "x").write_text("i-am-a-file")
+    (src / "x").mkdir()
+    (src / "x" / "inner.py").write_text("dir-content")
+
+    installer.move_tree(str(src), str(dst))
+
+    assert (dst / "x").is_dir()
+    assert (dst / "x" / "inner.py").read_text() == "dir-content"
+
+
+def test_exdev_file_copies_via_temp_then_atomic_replace(monkeypatch, tmp_path):
+    installer = load_installer()
+    src = tmp_path / "src"
+    dst = tmp_path / "dst"
+    src.mkdir()
+    dst.mkdir()
+    (dst / "f.py").write_text("OLD")
+    (src / "f.py").write_text("NEW")
+
+    real_replace = os.replace
+    seen = {"exdev": False, "part": False}
+
+    def exdev_for_direct_move(s, d):
+        if s == str(src / "f.py"):
+            seen["exdev"] = True
+            raise OSError(errno.EXDEV, "cross-device link")
+        if s.endswith(".part"):
+            seen["part"] = True
+        return real_replace(s, d)
+
+    monkeypatch.setattr(installer.os, "replace", exdev_for_direct_move)
+
+    installer.move_tree(str(src), str(dst))
+
+    assert (dst / "f.py").read_text() == "NEW"
+    assert seen["exdev"] and seen["part"]
+    # No leftover temp file.
+    assert not (dst / "f.py.part").exists()

packages/ai/python/tests/test_install_feature_resume.py

@@ -0,0 +1,238 @@
+"""Robustness tests for download_with_resume: the urllib fallback downloader.
+
+These simulate the network/disk failure modes users hit (flaky connections,
+Range-ignoring proxies, truncated bodies, disk full, dead URLs) with no real
+network, and assert the downloader either self-recovers or fails fast with a
+clear, actionable error instead of corrupting the archive or hanging.
+"""
+
+import builtins
+import errno
+import importlib.util
+import os
+import urllib.error
+
+import pytest
+
+
+def load_installer():
+    script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
+    spec = importlib.util.spec_from_file_location("install_feature_resume_under_test", script_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+class FakeResp:
+    """Minimal stand-in for the urlopen response context manager."""
+
+    def __init__(self, data: bytes, status: int = 200):
+        self._data = data
+        self._pos = 0
+        self.status = status
+
+    def read(self, n: int) -> bytes:
+        chunk = self._data[self._pos : self._pos + n]
+        self._pos += len(chunk)
+        return chunk
+
+    def getcode(self) -> int:
+        return self.status
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        return False
+
+
+@pytest.fixture(autouse=True)
+def no_sleep(monkeypatch):
+    """Never actually sleep during backoff; record calls instead."""
+    installer = load_installer()
+    calls = []
+    monkeypatch.setattr(installer.time, "sleep", lambda s: calls.append(s))
+    return calls
+
+
+def _patch_urlopen(monkeypatch, installer, handler):
+    """handler(req, call_index) -> FakeResp or raises."""
+    state = {"n": 0}
+
+    def fake_urlopen(req, timeout=None):
+        idx = state["n"]
+        state["n"] += 1
+        return handler(req, idx)
+
+    monkeypatch.setattr(installer.urllib.request, "urlopen", fake_urlopen)
+    return state
+
+
+def test_fresh_download_succeeds(monkeypatch, tmp_path):
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    data = b"x" * 4096
+    _patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
+
+    dest = tmp_path / "bundle.tar.gz"
+    installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
+
+    assert dest.read_bytes() == data
+    assert not (tmp_path / "bundle.tar.gz.partial").exists()
+    assert not (tmp_path / "bundle.tar.gz.meta").exists()
+
+
+def test_range_ignored_200_restarts_instead_of_corrupting(monkeypatch, tmp_path, no_sleep):
+    """A resumable partial exists, but the server ignores Range and returns 200
+    with the full body. The downloader must restart (truncate) rather than
+    append the full body onto the partial and corrupt the archive."""
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    data = b"GOOD" * 1024
+
+    dest = tmp_path / "bundle.tar.gz"
+    partial = tmp_path / "bundle.tar.gz.partial"
+    meta = tmp_path / "bundle.tar.gz.meta"
+    # A stale partial that would corrupt the file if appended to.
+    partial.write_bytes(b"STALEPARTIAL")
+    meta.write_text('{"bytesDownloaded": 12}')
+
+    # Server ignores the Range header and always sends the full body as 200.
+    _patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
+
+    installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
+
+    # Exactly the good bytes, not STALEPARTIAL + data.
+    assert dest.read_bytes() == data
+    assert no_sleep == []  # no retry needed; it restarted cleanly in one pass
+
+
+def test_proper_206_resume_appends(monkeypatch, tmp_path):
+    """When the server honors Range with 206, the partial is kept and only the
+    remaining bytes are appended."""
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    full = b"HEADER" + b"TAILDATA" * 512
+    head_len = 6
+
+    dest = tmp_path / "bundle.tar.gz"
+    partial = tmp_path / "bundle.tar.gz.partial"
+    meta = tmp_path / "bundle.tar.gz.meta"
+    partial.write_bytes(full[:head_len])
+    meta.write_text(f'{{"bytesDownloaded": {head_len}}}')
+
+    def handler(req, i):
+        assert req.get_header("Range") == f"bytes={head_len}-"
+        return FakeResp(full[head_len:], 206)
+
+    _patch_urlopen(monkeypatch, installer, handler)
+    installer.download_with_resume("https://h/f", str(dest), len(full), 2, 85)
+    assert dest.read_bytes() == full
+
+
+def test_disk_full_fails_fast_without_retry(monkeypatch, tmp_path, no_sleep):
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    data = b"y" * 4096
+    dest = tmp_path / "bundle.tar.gz"
+    partial_path = str(dest) + ".partial"
+
+    real_open = builtins.open
+
+    def fake_open(path, mode="r", *a, **k):
+        if str(path) == partial_path and ("w" in mode or "a" in mode):
+            real = real_open(path, mode, *a, **k)
+
+            class NoSpace:
+                def __enter__(self):
+                    return self
+
+                def __exit__(self, *_):
+                    real.close()
+                    return False
+
+                def write(self, _data):
+                    raise OSError(errno.ENOSPC, "No space left on device")
+
+            return NoSpace()
+        return real_open(path, mode, *a, **k)
+
+    monkeypatch.setattr(builtins, "open", fake_open)
+    _patch_urlopen(monkeypatch, installer, lambda req, i: FakeResp(data, 200))
+
+    with pytest.raises(RuntimeError, match="disk space"):
+        installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
+
+    assert no_sleep == []  # disk-full is not retried
+    assert not os.path.exists(partial_path)
+
+
+def test_http_404_fails_fast_without_retry(monkeypatch, tmp_path, no_sleep):
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    dest = tmp_path / "bundle.tar.gz"
+
+    def handler(req, i):
+        raise urllib.error.HTTPError("https://h/f", 404, "Not Found", {}, None)
+
+    _patch_urlopen(monkeypatch, installer, handler)
+
+    with pytest.raises(RuntimeError, match="HTTP 404"):
+        installer.download_with_resume("https://h/f", str(dest), 4096, 2, 85)
+
+    assert no_sleep == []  # a 404 won't fix on retry
+
+
+def test_429_is_retried(monkeypatch, tmp_path, no_sleep):
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    data = b"z" * 2048
+    dest = tmp_path / "bundle.tar.gz"
+
+    def handler(req, i):
+        if i == 0:
+            raise urllib.error.HTTPError("https://h/f", 429, "Too Many Requests", {}, None)
+        return FakeResp(data, 200)
+
+    _patch_urlopen(monkeypatch, installer, handler)
+    installer.download_with_resume("https://h/f", str(dest), len(data), 2, 85)
+
+    assert dest.read_bytes() == data
+    assert len(no_sleep) == 1  # backed off once, then succeeded
+
+
+def test_truncated_body_is_retried_then_succeeds(monkeypatch, tmp_path, no_sleep):
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    full = b"c" * 4096
+    dest = tmp_path / "bundle.tar.gz"
+
+    def handler(req, i):
+        if i == 0:
+            return FakeResp(full[:-100], 200)  # truncated / error page
+        return FakeResp(full, 200)
+
+    _patch_urlopen(monkeypatch, installer, handler)
+    installer.download_with_resume("https://h/f", str(dest), len(full), 2, 85)
+
+    assert dest.read_bytes() == full
+    assert len(no_sleep) == 1
+
+
+def test_connection_error_retried_then_raises_after_max(monkeypatch, tmp_path, no_sleep):
+    installer = load_installer()
+    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
+    dest = tmp_path / "bundle.tar.gz"
+
+    def handler(req, i):
+        raise urllib.error.URLError("connection reset by peer")
+
+    _patch_urlopen(monkeypatch, installer, handler)
+
+    with pytest.raises(RuntimeError, match="after 5 attempts"):
+        installer.download_with_resume("https://h/f", str(dest), 4096, 2, 85)
+
+    assert len(no_sleep) == 4  # 5 attempts => 4 backoffs
+    assert not os.path.exists(str(dest) + ".partial")
+    assert not os.path.exists(str(dest) + ".meta")