Merge pull request #96 from VibePod/loading-bar

Add loading bars when download agent, skills engine or datasette image

Author: nezhar

src/vibepod/core/docker.py

@@ -123,6 +123,18 @@ def _version_is_podman(version: Any) -> bool:
     return "podman" in str(version.get("Name", "")).lower()
 
 
+def _parse_image_name(image: str) -> tuple[str, str | None]:
+    """Parse a full image string into repository and tag/digest."""
+    if "@" in image:
+        repository, tag = image.split("@", 1)
+        return repository, tag
+    elif ":" in image:
+        parts = image.rsplit(":", 1)
+        if "/" not in parts[1]:
+            return parts[0], parts[1]
+    return image, None
+
+
 class DockerManager:
     """Manager for all Docker operations."""
 
@@ -162,12 +174,95 @@ def is_rootless_podman(self) -> bool:
         self._rootless_podman = rootless and _version_is_podman(version)
         return self._rootless_podman
 
-    def pull_image(self, image: str) -> None:
+    def _pull_image_with_progress(self, image: str) -> None:
+        from rich.progress import (
+            BarColumn,
+            DownloadColumn,
+            Progress,
+            TextColumn,
+            TimeRemainingColumn,
+            TransferSpeedColumn,
+        )
+
+        from vibepod.utils.console import console
+
+        repository, tag = _parse_image_name(image)
         try:
-            self.client.images.pull(image)
+            response = self.client.api.pull(repository, tag=tag, stream=True, decode=True)
         except APIError as exc:
             raise DockerClientError(f"Failed to pull image {image}: {exc}") from exc
 
+        tasks: dict[str, Any] = {}
+        try:
+            with Progress(
+                TextColumn("{task.description}"),
+                BarColumn(),
+                DownloadColumn(),
+                TransferSpeedColumn(),
+                TimeRemainingColumn(),
+                console=console,
+                disable=not console.is_terminal,
+            ) as progress:
+                for chunk in response:
+                    if not isinstance(chunk, dict):
+                        continue
+                    if "error" in chunk:
+                        raise DockerClientError(f"Failed to pull image {image}: {chunk['error']}")
+
+                    status = chunk.get("status", "")
+                    layer_id = chunk.get("id")
+                    progress_detail = chunk.get("progressDetail") or {}
+
+                    if not layer_id:
+                        if status and status not in (
+                            "Downloading",
+                            "Extracting",
+                            "Waiting",
+                            "Download complete",
+                            "Pull complete",
+                            "Already exists",
+                            "Pulling fs layer",
+                        ):
+                            progress.console.print(f"[dim]{status}[/dim]")
+                        continue
+
+                    status_color = "cyan"
+                    if status in ("Download complete", "Pull complete", "Already exists"):
+                        status_color = "green"
+                    elif status == "Waiting":
+                        status_color = "yellow"
+                    elif "error" in status.lower() or "fail" in status.lower():
+                        status_color = "red"
+
+                    description = f"[{status_color}][{layer_id}][/{status_color}] {status}"
+
+                    if layer_id not in tasks:
+                        tasks[layer_id] = progress.add_task(description, total=None)
+
+                    task_id = tasks[layer_id]
+                    progress.update(task_id, description=description)
+
+                    total = progress_detail.get("total", 0)
+                    current = progress_detail.get("current", 0)
+
+                    if total:
+                        progress.update(task_id, total=total, completed=current)
+
+                    if status in ("Download complete", "Pull complete", "Already exists"):
+                        if total:
+                            progress.update(task_id, completed=total, total=total)
+                        else:
+                            progress.update(task_id, completed=1, total=1)
+        except Exception as exc:
+            if isinstance(exc, DockerClientError):
+                raise
+            if isinstance(exc, APIError):
+                raise DockerClientError(f"Failed to pull image {image}: {exc}") from exc
+            raise DockerClientError(f"Failed to pull image {image}: {exc}") from exc
+
+    def pull_image(self, image: str) -> None:
+        self._pull_image_with_progress(image)
+
     def pull_if_newer(self, image: str) -> bool:
         """Pull *image* and return True if the local image was updated.
 
@@ -182,15 +277,15 @@ def pull_if_newer(self, image: str) -> bool:
             except NotFound:
                 old_id = None
 
-            self.client.images.pull(image)
+            self.pull_image(image)
 
             try:
                 new_id = self.client.images.get(image).id
             except NotFound:
                 return False
 
             return bool(old_id != new_id)
-        except APIError:
+        except (APIError, DockerClientError):
             return False
 
     def ensure_network(self, name: str) -> None:
@@ -443,6 +538,12 @@ def ensure_datasette(
                 return existing
             existing.remove(force=True)
 
+        if hasattr(self.client, "images"):
+            try:
+                self.client.images.get(image)
+            except NotFound:
+                self.pull_image(image)
+
         logs_db_path.parent.mkdir(parents=True, exist_ok=True)
         if not logs_db_path.exists():
             logs_db_path.touch()
@@ -489,6 +590,12 @@ def ensure_proxy(self, image: str, db_path: Path, ca_dir: Path, network: str) ->
                 return existing
             existing.remove(force=True)
 
+        if hasattr(self.client, "images"):
+            try:
+                self.client.images.get(image)
+            except NotFound:
+                self.pull_image(image)
+
         db_path.parent.mkdir(parents=True, exist_ok=True)
         ca_dir.mkdir(parents=True, exist_ok=True)
 

src/vibepod/core/skills_engine.py

@@ -16,9 +16,13 @@
     SKILLS_ENGINE_IMAGE,
     USER_SKILLS_DIR,
 )
+from vibepod.core.config import get_config
+from vibepod.core.docker import DockerClientError, DockerManager, NotFound
 
 Scope = Literal["local", "user"]
 
+_skills_engine_checked = False
+
 
 class SkillsEngineError(RuntimeError):
     """Raised when the driver cannot return an engine result."""
@@ -115,6 +119,39 @@ def run_engine(
     Returns the parsed JSON payload if ``json_output`` is True. Stderr from the
     engine (human-readable progress) is always captured but never parsed.
     """
+    global _skills_engine_checked
+    if not _skills_engine_checked:
+        try:
+            manager = DockerManager()
+            image_exists = False
+            try:
+                manager.client.images.get(SKILLS_ENGINE_IMAGE)
+                image_exists = True
+            except NotFound:
+                pass
+
+            config = get_config()
+            auto_pull_enabled = bool(config.get("auto_pull", True))
+            is_latest = (
+                ":" not in SKILLS_ENGINE_IMAGE.split("/")[-1]
+                or SKILLS_ENGINE_IMAGE.endswith(":latest")
+            )
+
+            if not image_exists:
+                manager.pull_image(SKILLS_ENGINE_IMAGE)
+            elif auto_pull_enabled and is_latest:
+                try:
+                    from vibepod.utils.console import info
+                    info("Checking for skills-engine image updates…")
+                    manager.pull_if_newer(SKILLS_ENGINE_IMAGE)
+                except Exception:
+                    pass
+            _skills_engine_checked = True
+        except DockerClientError as exc:
+            raise SkillsEngineError(str(exc)) from exc
+        except Exception as exc:
+            raise SkillsEngineError(f"Docker initialization failed: {exc}") from exc
+
     local, user, cache = _ensure_dirs(cwd, local_required=local_required)
 
     cmd: list[str] = [

tests/test_docker.py

@@ -0,0 +1,172 @@
+"""Tests for Docker manager image pulling and parsing helper functions."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vibepod.core.docker import (
+    APIError,
+    DockerClientError,
+    DockerManager,
+    NotFound,
+    _parse_image_name,
+)
+
+
+def test_parse_image_name() -> None:
+    assert _parse_image_name("vibepod/datasette:latest") == ("vibepod/datasette", "latest")
+    assert _parse_image_name("vibepod/datasette@sha256:abcd") == (
+        "vibepod/datasette",
+        "sha256:abcd",
+    )
+    assert _parse_image_name("localhost:5000/vibepod/datasette:latest") == (
+        "localhost:5000/vibepod/datasette",
+        "latest",
+    )
+    assert _parse_image_name("localhost:5000/vibepod/datasette") == (
+        "localhost:5000/vibepod/datasette",
+        None,
+    )
+    assert _parse_image_name("ubuntu") == ("ubuntu", None)
+
+
+@patch("vibepod.core.docker.docker")
+def test_pull_image_success(mock_docker) -> None:
+    # Set up mocks for DockerManager initialization
+    mock_client = MagicMock()
+    mock_docker.from_env.return_value = mock_client
+
+    # Define mock response from api.pull stream
+    mock_client.api.pull.return_value = [
+        {"status": "Pulling from library/ubuntu", "id": "latest"},
+        {"status": "Pulling fs layer", "id": "layer1"},
+        {"status": "Downloading", "id": "layer1", "progressDetail": {"current": 50, "total": 100}},
+        {"status": "Download complete", "id": "layer1"},
+        {"status": "Extracting", "id": "layer1", "progressDetail": {"current": 100, "total": 100}},
+        {"status": "Pull complete", "id": "layer1"},
+        {"status": "Already exists", "id": "layer2"},
+    ]
+
+    manager = DockerManager()
+    manager.pull_image("vibepod/datasette:latest")
+
+    mock_client.api.pull.assert_called_once_with(
+        "vibepod/datasette", tag="latest", stream=True, decode=True
+    )
+
+
+@patch("vibepod.core.docker.docker")
+def test_pull_image_api_error(mock_docker) -> None:
+    mock_client = MagicMock()
+    mock_docker.from_env.return_value = mock_client
+
+    # api.pull raises APIError
+    mock_client.api.pull.side_effect = APIError("Failed")
+
+    manager = DockerManager()
+    with pytest.raises(DockerClientError) as exc_info:
+        manager.pull_image("vibepod/datasette:latest")
+    assert "Failed to pull image" in str(exc_info.value)
+
+
+@patch("vibepod.core.docker.docker")
+def test_pull_image_chunk_error(mock_docker) -> None:
+    mock_client = MagicMock()
+    mock_docker.from_env.return_value = mock_client
+
+    # api.pull yields an error chunk
+    mock_client.api.pull.return_value = [
+        {"status": "Pulling fs layer", "id": "layer1"},
+        {"error": "Registry returned 404"},
+    ]
+
+    manager = DockerManager()
+    with pytest.raises(DockerClientError) as exc_info:
+        manager.pull_image("vibepod/datasette:latest")
+    assert "Registry returned 404" in str(exc_info.value)
+
+
+@patch("vibepod.core.docker.docker")
+def test_pull_if_newer(mock_docker) -> None:
+    mock_client = MagicMock()
+    mock_docker.from_env.return_value = mock_client
+
+    # Scenario 1: Image exists and gets updated (old_id != new_id)
+    mock_image_old = MagicMock()
+    mock_image_old.id = "sha256:old"
+    mock_image_new = MagicMock()
+    mock_image_new.id = "sha256:new"
+
+    mock_client.images.get.side_effect = [mock_image_old, mock_image_new]
+    mock_client.api.pull.return_value = [{"status": "Image is up to date"}]
+
+    manager = DockerManager()
+    updated = manager.pull_if_newer("vibepod/datasette:latest")
+    assert updated is True
+
+    # Scenario 2: Image is already up to date (old_id == new_id)
+    mock_client.images.get.side_effect = [mock_image_old, mock_image_old]
+    updated = manager.pull_if_newer("vibepod/datasette:latest")
+    assert updated is False
+
+    # Scenario 3: Image not found locally before, but pulled successfully
+    mock_client.images.get.side_effect = [NotFound("not found"), mock_image_new]
+    updated = manager.pull_if_newer("vibepod/datasette:latest")
+    assert updated is True
+
+    # Scenario 4: Pull fails
+    mock_client.images.get.side_effect = [NotFound("not found")]
+    mock_client.api.pull.side_effect = APIError("Failed")
+    updated = manager.pull_if_newer("vibepod/datasette:latest")
+    assert updated is False
+
+
+@patch("vibepod.core.docker.docker")
+def test_ensure_datasette_pulls_image_when_missing(mock_docker, tmp_path: Path) -> None:
+    mock_client = MagicMock()
+    mock_docker.from_env.return_value = mock_client
+
+    mock_client.containers.list.return_value = []
+    mock_client.images.get.side_effect = NotFound("Image not found")
+    mock_client.api.pull.return_value = [{"status": "Pulling"}]
+
+    manager = DockerManager()
+    manager.ensure_datasette(
+        image="vibepod/datasette:latest",
+        logs_db_path=tmp_path / "logs.db",
+        proxy_db_path=tmp_path / "proxy.db",
+        port=8001,
+    )
+
+    mock_client.api.pull.assert_called_once_with(
+        "vibepod/datasette", tag="latest", stream=True, decode=True
+    )
+    mock_client.containers.run.assert_called_once()
+
+
+@patch("vibepod.core.docker.docker")
+def test_ensure_proxy_pulls_image_when_missing(mock_docker, tmp_path: Path) -> None:
+    mock_client = MagicMock()
+    mock_docker.from_env.return_value = mock_client
+
+    mock_client.containers.list.return_value = []
+    mock_client.images.get.side_effect = NotFound("Image not found")
+    mock_client.api.pull.return_value = [{"status": "Pulling"}]
+
+    manager = DockerManager()
+    manager.is_rootless_podman = MagicMock(return_value=False)
+
+    manager.ensure_proxy(
+        image="vibepod/proxy:latest",
+        db_path=tmp_path / "proxy.db",
+        ca_dir=tmp_path / "ca",
+        network="vibepod-network",
+    )
+
+    mock_client.api.pull.assert_called_once_with(
+        "vibepod/proxy", tag="latest", stream=True, decode=True
+    )
+    mock_client.containers.run.assert_called_once()