Add Workspace.mount() API for programmatic branchfs setup

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

src/branching/core/workspace.py

@@ -1,11 +1,89 @@
 # SPDX-License-Identifier: Apache-2.0
 """Workspace class - main user-facing API for branch management."""
 
+from __future__ import annotations
+
+import shutil
+import subprocess
+import tempfile
+import time
 from pathlib import Path
+from typing import Optional
 
 from .base import FSBackend
 from .branch import Branch, OnSuccessAction, OnErrorAction
 from .registry import detect_fs_for_mount
+from ..exceptions import MountError
+
+
+class _OwnedMount:
+    """Manages a branchfs mount lifecycle started by Workspace.mount()."""
+
+    def __init__(
+        self,
+        base: Path,
+        mountpoint: Optional[Path],
+        storage: Optional[Path],
+        binary: Path,
+    ):
+        self.base = base
+        self.binary = binary
+        self._tmp_mountpoint = mountpoint is None
+        self._tmp_storage = storage is None
+        self.mountpoint = Path(mountpoint) if mountpoint else Path(
+            tempfile.mkdtemp(prefix="branchfs_mnt_")
+        )
+        self.storage = Path(storage) if storage else Path(
+            tempfile.mkdtemp(prefix="branchfs_storage_")
+        )
+
+    def start(self) -> None:
+        """Run `branchfs mount` which starts the daemon and mounts."""
+        self.mountpoint.mkdir(parents=True, exist_ok=True)
+        self.storage.mkdir(parents=True, exist_ok=True)
+
+        result = subprocess.run(
+            [
+                str(self.binary), "mount",
+                "--base", str(self.base),
+                "--storage", str(self.storage),
+                str(self.mountpoint),
+            ],
+            capture_output=True, text=True,
+        )
+        if result.returncode != 0:
+            self._cleanup_dirs()
+            raise MountError(
+                f"branchfs mount failed: {result.stderr.strip() or result.stdout.strip()}"
+            )
+
+        # Wait for FUSE to be ready
+        ctl = self.mountpoint / ".branchfs_ctl"
+        for _ in range(50):
+            if ctl.exists():
+                return
+            time.sleep(0.1)
+
+        self.stop()
+        raise MountError("branchfs mount timed out waiting for FUSE")
+
+    def stop(self) -> None:
+        """Unmount via branchfs CLI and clean up temp dirs."""
+        subprocess.run(
+            [
+                str(self.binary), "unmount",
+                str(self.mountpoint),
+                "--storage", str(self.storage),
+            ],
+            capture_output=True,
+        )
+        self._cleanup_dirs()
+
+    def _cleanup_dirs(self) -> None:
+        if self._tmp_mountpoint:
+            shutil.rmtree(self.mountpoint, ignore_errors=True)
+        if self._tmp_storage:
+            shutil.rmtree(self.storage, ignore_errors=True)
 
 
 class Workspace:
@@ -18,13 +96,14 @@ class Workspace:
     Example:
         from branching import Workspace
 
-        # Open workspace from existing mount
-        ws = Workspace("/mnt/main")
+        # Mount and use (handles setup and teardown)
+        with Workspace.mount("./my_project") as ws:
+            with ws.branch("attempt1") as b:
+                (b.path / "result.txt").write_text("done")
+                # auto-commits on success, auto-aborts on exception
 
-        # Simple speculative execution
-        with ws.branch("attempt1") as b:
-            subprocess.run(["agent", "--workdir", str(b.path)])
-            # auto-commits on success, auto-aborts on exception
+        # Or open an already-mounted workspace
+        ws = Workspace("/mnt/main")
 
     Attributes:
         path: Path to the main branch mountpoint
@@ -43,6 +122,71 @@ def __init__(self, path: str | Path):
         """
         self._path = Path(path).resolve()
         self._fs, self._mount_root = detect_fs_for_mount(self._path)
+        self._owned_mount: Optional[_OwnedMount] = None
+
+    @classmethod
+    def mount(
+        cls,
+        base: str | Path,
+        *,
+        mountpoint: str | Path | None = None,
+        storage: str | Path | None = None,
+        branchfs_bin: str | Path | None = None,
+    ) -> Workspace:
+        """
+        Mount a branchfs workspace and return a ready Workspace.
+
+        Handles daemon startup and FUSE mount automatically.
+        Use as a context manager for automatic cleanup::
+
+            with Workspace.mount("./project") as ws:
+                with ws.branch("fix") as b:
+                    ...
+
+        Or call ``ws.close()`` manually when done.
+
+        Args:
+            base: Directory to branch from (your project root).
+            mountpoint: Where to mount FUSE. Auto-created temp dir if None.
+            storage: Directory for daemon/branch data. Temp dir if None.
+            branchfs_bin: Path to branchfs binary. Searches PATH if None.
+
+        Raises:
+            MountError: If branchfs binary not found or mount fails.
+        """
+        base = Path(base).resolve()
+        if not base.is_dir():
+            raise MountError(f"Base directory does not exist: {base}")
+
+        binary = _find_branchfs(branchfs_bin)
+
+        owned = _OwnedMount(
+            base,
+            Path(mountpoint) if mountpoint else None,
+            Path(storage) if storage else None,
+            binary,
+        )
+        owned.start()
+
+        ws = cls(owned.mountpoint)
+        ws._owned_mount = owned
+        return ws
+
+    def close(self) -> None:
+        """Unmount and clean up if this workspace was created via mount().
+
+        No-op if the workspace was opened from an existing mount.
+        """
+        if self._owned_mount is not None:
+            self._owned_mount.stop()
+            self._owned_mount = None
+
+    def __enter__(self) -> Workspace:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        self.close()
+        return False
 
     @property
     def path(self) -> Path:
@@ -97,3 +241,21 @@ def branch(
 
     def __repr__(self) -> str:
         return f"Workspace(path={self._path!r}, fstype={self.fstype!r})"
+
+
+def _find_branchfs(hint: str | Path | None) -> Path:
+    """Locate the branchfs binary."""
+    if hint is not None:
+        p = Path(hint)
+        if p.is_file():
+            return p
+        raise MountError(f"branchfs binary not found at {p}")
+
+    found = shutil.which("branchfs")
+    if found:
+        return Path(found)
+
+    raise MountError(
+        "branchfs binary not found in PATH. "
+        "Install branchfs or pass branchfs_bin= to Workspace.mount()."
+    )

tests/test_workspace_mount.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Integration tests for Workspace.mount() — requires branchfs binary and FUSE."""
+
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from branching.core.workspace import Workspace, _find_branchfs
+from branching.exceptions import MountError
+
+
+# Skip all tests if branchfs binary not available or no FUSE
+pytestmark = pytest.mark.skipif(
+    shutil.which("branchfs") is None
+    or not (Path("/dev/fuse").exists() and os.access("/dev/fuse", os.R_OK | os.W_OK)),
+    reason="branchfs binary not in PATH or /dev/fuse not accessible",
+)
+
+
+@pytest.fixture
+def base_dir():
+    d = Path(tempfile.mkdtemp(prefix="branchfs_test_base_"))
+    (d / "file1.txt").write_text("base content")
+    (d / "subdir").mkdir()
+    (d / "subdir" / "nested.txt").write_text("nested")
+    yield d
+    shutil.rmtree(d, ignore_errors=True)
+
+
+class TestWorkspaceMount:
+    def test_mount_and_close(self, base_dir):
+        ws = Workspace.mount(base_dir)
+        try:
+            assert ws.path.exists()
+            assert (ws.path / ".branchfs_ctl").exists()
+            assert (ws.path / "file1.txt").read_text() == "base content"
+            assert ws.fstype == "fuse.branchfs"
+        finally:
+            ws.close()
+
+    def test_context_manager(self, base_dir):
+        with Workspace.mount(base_dir) as ws:
+            assert (ws.path / "file1.txt").read_text() == "base content"
+        # After exit, mount should be cleaned up
+        assert not (ws.path / ".branchfs_ctl").exists()
+
+    def test_branch_commit(self, base_dir):
+        with Workspace.mount(base_dir) as ws:
+            with ws.branch("test_branch") as b:
+                (b.path / "new_file.txt").write_text("from branch")
+            # auto-committed
+            assert (base_dir / "new_file.txt").read_text() == "from branch"
+
+    def test_branch_abort(self, base_dir):
+        with Workspace.mount(base_dir) as ws:
+            try:
+                with ws.branch("fail_branch") as b:
+                    (b.path / "bad_file.txt").write_text("should vanish")
+                    raise ValueError("simulated failure")
+            except ValueError:
+                pass
+            # auto-aborted — file should not be in base
+            assert not (base_dir / "bad_file.txt").exists()
+
+    def test_explicit_mountpoint_and_storage(self, base_dir, tmp_path):
+        mnt = tmp_path / "mnt"
+        stor = tmp_path / "stor"
+        with Workspace.mount(base_dir, mountpoint=mnt, storage=stor) as ws:
+            assert ws.path == mnt.resolve()
+            assert (ws.path / "file1.txt").exists()
+
+    def test_mount_nonexistent_base_raises(self):
+        with pytest.raises(MountError, match="does not exist"):
+            Workspace.mount("/nonexistent/path/abc123")
+
+    def test_mount_bad_binary_raises(self, base_dir):
+        with pytest.raises(MountError, match="not found"):
+            Workspace.mount(base_dir, branchfs_bin="/no/such/binary")