fix: persist packaged SEA binaries during extraction

src/stagehand/_custom/sea_binary.py

@@ -43,7 +43,7 @@ def _ensure_executable(path: Path) -> None:
         path.chmod(mode | 0o100)
 
 
-def _resource_binary_path(filename: str) -> Path | None:
+def _resource_binary_path(filename: str, *, version: str) -> Path | None:
     # Expect binaries to live at stagehand/_sea/<filename> inside the installed package.
     try:
         root = importlib_resources.files("stagehand")
@@ -58,7 +58,9 @@ def _resource_binary_path(filename: str) -> Path | None:
         return None
 
     with importlib_resources.as_file(candidate) as extracted:
-        return extracted
+        # ZIP-backed resources are temporary and disappear when this context
+        # exits, so persist the binary before returning its path.
+        return _copy_to_cache(src=extracted, filename=filename, version=version)
 
 
 def _copy_to_cache(*, src: Path, filename: str, version: str) -> Path:
@@ -97,12 +99,11 @@ def resolve_binary_path(
     filename = default_binary_filename()
 
     # Prefer packaged resources (works for wheel installs).
-    resource_path = _resource_binary_path(filename)
+    if version is None:
+        version = os.environ.get("STAGEHAND_VERSION") or __version__
+    resource_path = _resource_binary_path(filename, version=version)
     if resource_path is not None:
-        # Best-effort versioning to keep cached binaries stable across upgrades.
-        if version is None:
-            version = os.environ.get("STAGEHAND_VERSION") or __version__
-        return _copy_to_cache(src=resource_path, filename=filename, version=version)
+        return resource_path
 
     # Fallback: source checkout layout (works for local dev in-repo).
     here = Path(__file__).resolve()

tests/test_sea_binary.py

@@ -2,6 +2,8 @@
 
 import importlib.util
 from pathlib import Path
+from contextlib import contextmanager
+from collections.abc import Iterator
 
 import pytest
 
@@ -37,27 +39,55 @@ def test_resolve_binary_path_defaults_cache_version_to_package_version(
     monkeypatch.delenv("STAGEHAND_SEA_BINARY", raising=False)
     monkeypatch.delenv("STAGEHAND_VERSION", raising=False)
 
-    def _fake_resource_binary_path(_filename: str) -> Path:
+    def _fake_resource_binary_path(_filename: str, *, version: str) -> Path:
+        captured["version"] = version
         return resource_path
 
     monkeypatch.setattr(sea_binary, "_resource_binary_path", _fake_resource_binary_path)
 
-    def _fake_copy_to_cache(*, src: Path, filename: str, version: str) -> Path:
-        captured["src"] = src
-        captured["filename"] = filename
-        captured["version"] = version
-        return tmp_path / "cache" / filename
-
-    monkeypatch.setattr(sea_binary, "_copy_to_cache", _fake_copy_to_cache)
-
     resolved = sea_binary.resolve_binary_path()
 
-    assert resolved == tmp_path / "cache" / sea_binary.default_binary_filename()
-    assert captured["src"] == resource_path
-    assert captured["filename"] == sea_binary.default_binary_filename()
+    assert resolved == resource_path
     assert captured["version"] == __version__
 
 
+def test_resolve_binary_path_caches_resource_before_context_exits(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    class FakeResource:
+        def joinpath(self, _name: str) -> FakeResource:
+            return self
+
+        def is_file(self) -> bool:
+            return True
+
+    extracted = tmp_path / "temporary-binary"
+
+    @contextmanager
+    def fake_as_file(_resource: object) -> Iterator[Path]:
+        extracted.write_bytes(b"binary")
+        try:
+            yield extracted
+        finally:
+            extracted.unlink()
+
+    def fake_files(_package: str) -> FakeResource:
+        return FakeResource()
+
+    monkeypatch.delenv("STAGEHAND_SEA_BINARY", raising=False)
+    monkeypatch.setattr(sea_binary.importlib_resources, "files", fake_files)
+    monkeypatch.setattr(sea_binary.importlib_resources, "as_file", fake_as_file)
+    monkeypatch.setattr(sea_binary, "_cache_dir", lambda: tmp_path / "cache")
+    monkeypatch.setattr(sea_binary, "default_binary_filename", lambda: "stagehand-test")
+
+    resolved = sea_binary.resolve_binary_path(version="test")
+
+    assert resolved == tmp_path / "cache" / "test" / "stagehand-test"
+    assert resolved.read_bytes() == b"binary"
+    assert not extracted.exists()
+
+
 def test_parse_server_tag_rejects_prerelease_tags() -> None:
     assert download_binary._parse_server_tag("stagehand-server-v3/v3.20.0-dev") is None
     assert download_binary._parse_server_tag("stagehand-server-v3/v3.20.0+build.1") is None