fix: make SEA binary cache writes concurrency-safe

Author: abhinavkr26104

src/stagehand/_custom/sea_binary.py

@@ -2,8 +2,8 @@
 
 import os
 import sys
-import hashlib
 import platform
+import tempfile
 import importlib.resources as importlib_resources
 from pathlib import Path
 from contextlib import suppress
@@ -71,9 +71,21 @@ def _copy_to_cache(*, src: Path, filename: str, version: str) -> Path:
         return dst
 
     data = src.read_bytes()
-    tmp = cache_root / f".{filename}.{hashlib.sha256(data).hexdigest()}.tmp"
-    tmp.write_bytes(data)
-    tmp.replace(dst)
+    with tempfile.NamedTemporaryFile(dir=cache_root, prefix=f".{filename}.", suffix=".tmp", delete=False) as file:
+        file.write(data)
+        tmp = Path(file.name)
+
+    try:
+        try:
+            tmp.replace(dst)
+        except OSError:
+            # Another process may have populated the cache first. Its atomic
+            # replace guarantees that an existing destination is complete.
+            if not dst.exists():
+                raise
+    finally:
+        tmp.unlink(missing_ok=True)
+
     _ensure_executable(dst)
     return dst
 

tests/test_sea_binary.py

@@ -2,6 +2,7 @@
 
 import importlib.util
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
 
 import pytest
 
@@ -23,6 +24,37 @@ def _load_download_binary_module():
 download_binary = _load_download_binary_module()
 
 
+def test_copy_to_cache_reuses_existing_binary(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    source = tmp_path / "source"
+    source.write_bytes(b"new")
+    cached = tmp_path / "cache" / "test" / "stagehand-test"
+    cached.parent.mkdir(parents=True)
+    cached.write_bytes(b"cached")
+    monkeypatch.setattr(sea_binary, "_cache_dir", lambda: tmp_path / "cache")
+
+    result = sea_binary._copy_to_cache(src=source, filename="stagehand-test", version="test")
+
+    assert result == cached
+    assert result.read_bytes() == b"cached"
+
+
+def test_copy_to_cache_is_safe_when_called_concurrently(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    source = tmp_path / "source"
+    source.write_bytes(b"binary" * 1_000_000)
+    monkeypatch.setattr(sea_binary, "_cache_dir", lambda: tmp_path / "cache")
+
+    def copy(_index: int) -> Path:
+        return sea_binary._copy_to_cache(src=source, filename="stagehand-test", version="test")
+
+    with ThreadPoolExecutor(max_workers=16) as executor:
+        results = list(executor.map(copy, range(16)))
+
+    expected = tmp_path / "cache" / "test" / "stagehand-test"
+    assert results == [expected] * 16
+    assert expected.read_bytes() == source.read_bytes()
+    assert list(expected.parent.glob("*.tmp")) == []
+
+
 def test_resolve_binary_path_defaults_cache_version_to_package_version(
     monkeypatch: pytest.MonkeyPatch,
     tmp_path: Path,