feat: extract lfs if the file contents has changed

dimos/utils/data.py

@@ -14,17 +14,21 @@
 
 from datetime import datetime
 from functools import cache
+import json
 import os
 from pathlib import Path
 import platform
 import re
+import shutil
 import subprocess
 import sys
 import tarfile
 import tempfile
 import time
 
-from dimos.constants import DIMOS_PROJECT_ROOT
+from filelock import FileLock
+
+from dimos.constants import DIMOS_PROJECT_ROOT, STATE_DIR
 from dimos.utils.logging_config import setup_logger
 
 logger = setup_logger()
@@ -221,9 +225,89 @@ def _lfs_pull(file_path: Path, repo_root: Path, *, retries: int = 2) -> None:
 def _decompress_archive(filename: str | Path) -> Path:
     target_dir = get_data_dir()
     filename_path = Path(filename)
-    with tarfile.open(filename_path, "r:gz") as tar:
-        tar.extractall(target_dir)
-    return target_dir / filename_path.name.replace(".tar.gz", "")
+    extracted_name = filename_path.name.replace(".tar.gz", "")
+    extracted_path = target_dir / extracted_name
+
+    # extract the file to a temp location first, then swap if that worked
+    staging_dir = Path(tempfile.mkdtemp(dir=target_dir, prefix=f".{extracted_name}.staging-"))
+    try:
+        with tarfile.open(filename_path, "r:gz") as tar:
+            tar.extractall(staging_dir)
+
+        # validate before touching the old extraction, so a mis-packaged
+        # archive can't delete a good copy and then fail to replace it
+        staged_path = staging_dir / extracted_name
+        if not staged_path.exists():
+            raise RuntimeError(
+                f"Archive '{filename_path.name}' has no top-level '{extracted_name}' member"
+            )
+
+        if extracted_path.is_dir():
+            shutil.rmtree(extracted_path)
+        elif extracted_path.exists():
+            extracted_path.unlink()
+        os.replace(staged_path, extracted_path)
+    finally:
+        shutil.rmtree(staging_dir, ignore_errors=True)
+
+    return extracted_path
+
+
+def _lfs_archive_path(archive_name: str | Path) -> Path:
+    return _get_lfs_dir() / (str(archive_name) + ".tar.gz")
+
+
+@cache
+def _lfs_extraction_state_dir() -> Path:
+    """Per-checkout directory for extraction stamps and locks."""
+    # Keyed by the data dir's own path so worktrees don't collide.
+    checkout_key = str(get_data_dir().resolve()).replace("/", "_").lstrip("_")
+    return STATE_DIR / "lfs-extractions" / checkout_key
+
+
+def _extraction_stamp_path(archive_name: str | Path) -> Path:
+    return _lfs_extraction_state_dir() / f"{archive_name}.json"
+
+
+def _archive_lock_path(archive_name: str | Path) -> Path:
+    return _lfs_extraction_state_dir() / f"{archive_name}.lock"
+
+
+def _archive_lock(archive_name: str | Path) -> FileLock:
+    _lfs_extraction_state_dir().mkdir(parents=True, exist_ok=True)
+    return FileLock(_archive_lock_path(archive_name))
+
+
+def _archive_identity(archive_path: Path) -> dict[str, int]:
+    stat_result = archive_path.stat()
+    return {"st_size": stat_result.st_size, "st_mtime_ns": stat_result.st_mtime_ns}
+
+
+def _extraction_is_current(archive_name: str | Path) -> bool:
+    """Whether the extracted data still matches its archive.
+
+    No archive behind it means a local recording: always current.
+    """
+    archive_path = _lfs_archive_path(archive_name)
+    if not archive_path.exists():
+        return True
+
+    stamp_path = _extraction_stamp_path(archive_name)
+    if not stamp_path.exists():
+        return False
+
+    try:
+        stamp: dict[str, int] = json.loads(stamp_path.read_text())
+        return stamp == _archive_identity(archive_path)
+    except (json.JSONDecodeError, OSError):
+        # if we can't load the stamp, just treat it as stale
+        return False
+
+
+def _write_extraction_stamp(archive_name: str | Path, archive_path: Path) -> None:
+    stamp_path = _extraction_stamp_path(archive_name)
+    stamp_path.parent.mkdir(parents=True, exist_ok=True)
+    stamp_path.write_text(json.dumps(_archive_identity(archive_path)))
 
 
 def _pull_lfs_archive(filename: str | Path) -> Path:
@@ -234,7 +318,7 @@ def _pull_lfs_archive(filename: str | Path) -> Path:
     repo_root = get_project_root()
 
     # Construct path to test data file
-    file_path = _get_lfs_dir() / (str(filename) + ".tar.gz")
+    file_path = _lfs_archive_path(filename)
 
     # Check if file exists
     if not file_path.exists():
@@ -291,17 +375,25 @@ def get_data(name: str | Path) -> Path:
     data_dir = get_data_dir()
     file_path = data_dir / name
 
-    # already pulled and decompressed, return it directly
-    if file_path.exists():
-        return file_path
-
     # extract archive root (first path component) and nested path
     path_parts = Path(name).parts
     archive_name = path_parts[0]
     nested_path = Path(*path_parts[1:]) if len(path_parts) > 1 else None
 
-    # download and decompress the archive root
-    archive_path = _decompress_archive(_pull_lfs_archive(archive_name))
+    def is_ready() -> bool:
+        # already extracted and current
+        return file_path.exists() and _extraction_is_current(archive_name)
+
+    if is_ready():
+        return file_path
+
+    with _archive_lock(archive_name):
+        if not is_ready():
+            archive_tar_path = _pull_lfs_archive(archive_name)
+            archive_path = _decompress_archive(archive_tar_path)
+            _write_extraction_stamp(archive_name, archive_tar_path)
+        else:
+            archive_path = data_dir / archive_name
 
     # return full path including nested components
     if nested_path:

dimos/utils/test_data.py

@@ -12,10 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
 import hashlib
+import io
+import json
 import os
 from pathlib import Path
 import subprocess
+import tarfile
 
 import pytest
 
@@ -100,6 +106,222 @@ def test_backup_file_keep_last_zero_removes_all(tmp_path: Path) -> None:
     assert list(tmp_path.glob("recording_go2.*.db")) == []
 
 
+@dataclass
+class DataLayout:
+    data_dir: Path
+    lfs_dir: Path
+    state_dir: Path
+    staging_root: Path
+
+
+@pytest.fixture
+def data_layout(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> DataLayout:
+    layout = DataLayout(
+        data_dir=tmp_path / "data",
+        lfs_dir=tmp_path / "data" / ".lfs",
+        state_dir=tmp_path / "state",
+        staging_root=tmp_path / "staging",
+    )
+    layout.data_dir.mkdir()
+    layout.lfs_dir.mkdir()
+    layout.staging_root.mkdir()
+
+    monkeypatch.setattr(
+        data,
+        "get_data_dir",
+        lambda extra_path=None: layout.data_dir / extra_path if extra_path else layout.data_dir,
+    )
+    monkeypatch.setattr(data, "_get_lfs_dir", lambda: layout.lfs_dir)
+    monkeypatch.setattr(data, "_lfs_extraction_state_dir", lambda: layout.state_dir)
+    monkeypatch.setattr(data, "_check_git_lfs_available", lambda: True)
+    return layout
+
+
+def _write_tar_gz(tar_path: Path, archive_name: str, content: bytes, staging_root: Path) -> None:
+    staging = staging_root / archive_name
+    staging.mkdir(exist_ok=True)
+    (staging / "payload.bin").write_bytes(content)
+    with tarfile.open(tar_path, "w:gz") as tar:
+        tar.add(staging, arcname=archive_name)
+
+
+def _write_single_file_tar_gz(tar_path: Path, member_name: str, content: bytes) -> None:
+    with tarfile.open(tar_path, "w:gz") as tar:
+        info = tarfile.TarInfo(member_name)
+        info.size = len(content)
+        tar.addfile(info, io.BytesIO(content))
+
+
+def _counting_decompress(
+    monkeypatch: pytest.MonkeyPatch,
+) -> Callable[[], int]:
+    """Wrap _decompress_archive and return its running call count."""
+    call_count = 0
+    original_decompress = data._decompress_archive
+
+    def counting_decompress(filename: str | Path) -> Path:
+        nonlocal call_count
+        call_count += 1
+        return original_decompress(filename)
+
+    monkeypatch.setattr(data, "_decompress_archive", counting_decompress)
+    return lambda: call_count
+
+
+def test_get_data_self_heals_legacy_extraction_without_stamp(data_layout: DataLayout) -> None:
+    """No stamp file means not trusted: re-extracts once, then self-heals."""
+    tar_path = data._lfs_archive_path("dataset")
+    _write_tar_gz(tar_path, "dataset", b"current", data_layout.staging_root)
+
+    # archive and extraction both present, but no stamp
+    (data_layout.data_dir / "dataset").mkdir()
+    (data_layout.data_dir / "dataset" / "payload.bin").write_bytes(b"stale")
+
+    healed = data.get_data("dataset/payload.bin")
+    assert healed.read_bytes() == b"current"
+    assert data._extraction_stamp_path("dataset").exists()
+
+
+def test_get_data_returns_local_recording_without_backing_archive(
+    data_layout: DataLayout, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    local_db = data_layout.data_dir / "recording.db"
+    local_db.write_bytes(b"local only, no archive behind it")
+
+    def fail_pull(name: str) -> Path:
+        raise AssertionError("must not attempt an LFS pull for a local-only recording")
+
+    monkeypatch.setattr(data, "_pull_lfs_archive", fail_pull)
+
+    assert data.get_data("recording.db") == local_db
+
+
+def test_get_data_reextracts_when_archive_changes_then_stabilizes(
+    data_layout: DataLayout, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    tar_path = data._lfs_archive_path("dataset")
+    _write_tar_gz(tar_path, "dataset", b"v1", data_layout.staging_root)
+
+    first = data.get_data("dataset/payload.bin")
+    assert first.read_bytes() == b"v1"
+
+    call_count = _counting_decompress(monkeypatch)
+
+    # unchanged archive: no re-extraction
+    data.get_data("dataset/payload.bin")
+    assert call_count() == 0
+
+    # archive re-tarred under the same name with different content
+    _write_tar_gz(
+        tar_path, "dataset", b"v2, a longer payload than before", data_layout.staging_root
+    )
+
+    second = data.get_data("dataset/payload.bin")
+    assert second.read_bytes() == b"v2, a longer payload than before"
+    assert call_count() == 1
+
+    # settles back to no re-extraction once the stamp matches again
+    data.get_data("dataset/payload.bin")
+    assert call_count() == 1
+
+
+def test_get_data_reextraction_drops_members_removed_from_archive(data_layout: DataLayout) -> None:
+    """Members dropped from a re-tarred archive must not survive on disk."""
+    staging = data_layout.staging_root / "dataset"
+    staging.mkdir(parents=True)
+    tar_path = data._lfs_archive_path("dataset")
+
+    (staging / "keep.bin").write_bytes(b"kept")
+    (staging / "drop.bin").write_bytes(b"dropped")
+    with tarfile.open(tar_path, "w:gz") as tar:
+        tar.add(staging, arcname="dataset")
+
+    dropped_path = data.get_data("dataset/drop.bin")
+    assert dropped_path.exists()
+
+    # re-tar without drop.bin
+    (staging / "drop.bin").unlink()
+    with tarfile.open(tar_path, "w:gz") as tar:
+        tar.add(staging, arcname="dataset")
+
+    kept_path = data.get_data("dataset/keep.bin")
+    assert kept_path.read_bytes() == b"kept"
+    assert not dropped_path.exists()
+
+
+def test_get_data_mispackaged_archive_does_not_destroy_prior_extraction(
+    data_layout: DataLayout,
+) -> None:
+    """A stale-triggering archive missing its top-level member must raise
+    without deleting the still-good previous extraction."""
+    tar_path = data._lfs_archive_path("dataset")
+    _write_tar_gz(tar_path, "dataset", b"v1", data_layout.staging_root)
+
+    first = data.get_data("dataset/payload.bin")
+    assert first.read_bytes() == b"v1"
+
+    # re-tar under the wrong top-level name, so extraction succeeds but the
+    # expected "dataset" member never lands in the staging dir
+    with tarfile.open(tar_path, "w:gz") as tar:
+        info = tarfile.TarInfo("wrong_name/payload.bin")
+        payload = b"v2"
+        info.size = len(payload)
+        tar.addfile(info, io.BytesIO(payload))
+
+    with pytest.raises(RuntimeError, match="dataset"):
+        data.get_data("dataset/payload.bin")
+
+    assert first.exists()
+    assert first.read_bytes() == b"v1"
+
+
+def test_get_data_reextracts_when_stamp_is_corrupt(data_layout: DataLayout) -> None:
+    tar_path = data._lfs_archive_path("dataset")
+    _write_tar_gz(tar_path, "dataset", b"current", data_layout.staging_root)
+
+    stamp_path = data._extraction_stamp_path("dataset")
+    stamp_path.parent.mkdir(parents=True, exist_ok=True)
+    stamp_path.write_text("not valid json")
+    (data_layout.data_dir / "dataset").mkdir()
+    (data_layout.data_dir / "dataset" / "payload.bin").write_bytes(b"stale")
+
+    result = data.get_data("dataset/payload.bin")
+
+    assert result.read_bytes() == b"current"
+    assert json.loads(stamp_path.read_text())
+
+
+def test_get_data_reextracts_single_file_archive(data_layout: DataLayout) -> None:
+    """A prior extraction that is a bare file, not a directory, must also be replaced."""
+    tar_path = data._lfs_archive_path("cafe.jpg")
+    _write_single_file_tar_gz(tar_path, "cafe.jpg", b"v1")
+
+    first = data.get_data("cafe.jpg")
+    assert first.is_file()
+    assert first.read_bytes() == b"v1"
+
+    _write_single_file_tar_gz(tar_path, "cafe.jpg", b"v2, a longer payload than before")
+
+    second = data.get_data("cafe.jpg")
+    assert second.read_bytes() == b"v2, a longer payload than before"
+
+
+def test_get_data_serializes_concurrent_extraction_of_same_archive(
+    data_layout: DataLayout, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Concurrent callers must not race: only one extraction happens."""
+    tar_path = data._lfs_archive_path("dataset")
+    _write_tar_gz(tar_path, "dataset", b"v1", data_layout.staging_root)
+
+    call_count = _counting_decompress(monkeypatch)
+
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        results = list(executor.map(lambda _: data.get_data("dataset/payload.bin"), range(8)))
+
+    assert all(result.read_bytes() == b"v1" for result in results)
+    assert call_count() == 1
+
+
 @pytest.mark.self_hosted
 def test_pull_file() -> None:
     repo_root = data.get_project_root()