feat(snapshots): wait for a snapshot to be Ready

.changeset/snapshot-wait.md

@@ -0,0 +1,5 @@
+---
+"neevai": minor
+---
+
+Wait for a snapshot to finish. `client.sandboxes.wait_for_snapshot(snapshot_id, timeout_ms=300000, poll_interval_ms=2000)` polls until the snapshot is Ready and returns it, and raises `NeevAIError` if it fails (with its reason) or is not Ready in time. `sandbox.snapshot(params, wait_until_ready=True)` takes a snapshot and waits the same way. Both work on the sync and async clients, as the JS SDK's `waitForSnapshot` and `snapshot({ waitUntilReady: true })` do.

src/neevai/handles/sandbox.py

@@ -10,6 +10,9 @@
 from neevai.resources.sandboxes import (
     DEFAULT_PORT_POLL_INTERVAL_MS,
     DEFAULT_PORT_WAIT_TIMEOUT_MS,
+    DEFAULT_SNAPSHOT_POLL_INTERVAL_MS,
+    DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS,
+    check_snapshot_wait,
 )
 from neevai.types import (
     AuditTrail,
@@ -368,15 +371,34 @@ def audit(
     def snapshot(
         self,
         params: CreateSnapshotParams | Mapping[str, Any] | None = None,
+        *,
+        wait_until_ready: bool = False,
+        timeout_ms: int = DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS,
+        poll_interval_ms: int = DEFAULT_SNAPSHOT_POLL_INTERVAL_MS,
     ) -> Snapshot:
-        """Captures a snapshot of this sandbox (returns immediately with status Pending)."""
+        """Captures a snapshot of this sandbox.
+
+        Returns at once with status Pending; with ``wait_until_ready=True`` it waits
+        and returns the Ready snapshot, as ``wait_for_snapshot`` does.
+        """
         if self.sandboxes is None:
             raise NeevAIError("Cannot snapshot a sandbox handle with no client context.")
-        return self.sandboxes.create_snapshot(
-            self.id,
-            params,
-            org_id=self.scope.org_id if self.scope else None,
-            project_id=self.scope.project_id if self.scope else None,
+        if wait_until_ready:
+            # Reject a bad timing before the snapshot exists, so none is left behind.
+            check_snapshot_wait(timeout_ms, poll_interval_ms)
+        org_id = self.scope.org_id if self.scope else None
+        project_id = self.scope.project_id if self.scope else None
+        snapshot = self.sandboxes.create_snapshot(
+            self.id, params, org_id=org_id, project_id=project_id
+        )
+        if not wait_until_ready:
+            return snapshot
+        return self.sandboxes.wait_for_snapshot(
+            snapshot.id,
+            timeout_ms=timeout_ms,
+            poll_interval_ms=poll_interval_ms,
+            org_id=org_id,
+            project_id=project_id,
         )
 
     def snapshots(self) -> list[Snapshot]:
@@ -868,14 +890,30 @@ async def audit(
     async def snapshot(
         self,
         params: CreateSnapshotParams | Mapping[str, Any] | None = None,
+        *,
+        wait_until_ready: bool = False,
+        timeout_ms: int = DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS,
+        poll_interval_ms: int = DEFAULT_SNAPSHOT_POLL_INTERVAL_MS,
     ) -> Snapshot:
+        """Captures a snapshot asynchronously; ``wait_until_ready=True`` waits for Ready."""
         if self.sandboxes is None:
             raise NeevAIError("Cannot snapshot a sandbox handle with no client context.")
-        return await self.sandboxes.create_snapshot(
-            self.id,
-            params,
-            org_id=self.scope.org_id if self.scope else None,
-            project_id=self.scope.project_id if self.scope else None,
+        if wait_until_ready:
+            # Reject a bad timing before the snapshot exists, so none is left behind.
+            check_snapshot_wait(timeout_ms, poll_interval_ms)
+        org_id = self.scope.org_id if self.scope else None
+        project_id = self.scope.project_id if self.scope else None
+        snapshot = await self.sandboxes.create_snapshot(
+            self.id, params, org_id=org_id, project_id=project_id
+        )
+        if not wait_until_ready:
+            return snapshot
+        return await self.sandboxes.wait_for_snapshot(
+            snapshot.id,
+            timeout_ms=timeout_ms,
+            poll_interval_ms=poll_interval_ms,
+            org_id=org_id,
+            project_id=project_id,
         )
 
     async def snapshots(self) -> list[Snapshot]:

src/neevai/resources/sandboxes.py

@@ -201,6 +201,33 @@ def _list_query(
     return query
 
 
+# Capturing a snapshot can take minutes, so it waits longer than a sandbox readiness wait.
+DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS = 300_000
+DEFAULT_SNAPSHOT_POLL_INTERVAL_MS = 2_000
+
+
+def check_snapshot_wait(timeout_ms: float, poll_interval_ms: float) -> None:
+    """Raises NeevAIError unless both snapshot-wait timings are positive and finite."""
+    # A zero, negative or NaN timing would spin a near-zero-delay poll or never time out.
+    for name, value in (("timeout_ms", timeout_ms), ("poll_interval_ms", poll_interval_ms)):
+        if not value > 0 or value == float("inf"):
+            raise NeevAIError(f"wait_for_snapshot: {name} must be a positive number (got {value}).")
+
+
+def _snapshot_done(snapshot: Snapshot, snapshot_id: str | UUID) -> bool:
+    if snapshot.status == "Ready":
+        return True
+    if snapshot.status == "Failed":
+        reason = f": {snapshot.error_message}" if snapshot.error_message else ""
+        raise NeevAIError(f"Snapshot {snapshot_id} failed to capture{reason}.")
+    return False
+
+
+def _snapshot_timeout(snapshot: Snapshot, snapshot_id: str | UUID, timeout_ms: float) -> str:
+    status = getattr(snapshot.status, "value", snapshot.status)
+    return f"Snapshot {snapshot_id} was not Ready within {timeout_ms}ms (status: {status})."
+
+
 class Sandboxes:
     """Operations on the /sandboxes API endpoint (synchronous)."""
 
@@ -608,6 +635,31 @@ def get_snapshot(
         raw = self._client._transport.request("GET", path)
         return coerce_model(Snapshot, raw)
 
+    def wait_for_snapshot(
+        self,
+        snapshot_id: str | UUID,
+        timeout_ms: int = DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS,
+        poll_interval_ms: int = DEFAULT_SNAPSHOT_POLL_INTERVAL_MS,
+        org_id: str | None = None,
+        project_id: str | None = None,
+    ) -> Snapshot:
+        """Polls a snapshot until it is Ready and returns it.
+
+        Raises:
+            NeevAIError: the snapshot Failed, was not Ready within ``timeout_ms``, or a
+                timing is not a positive number.
+        """
+        check_snapshot_wait(timeout_ms, poll_interval_ms)
+        deadline = time.monotonic() * 1000.0 + timeout_ms
+        while True:
+            snapshot = self.get_snapshot(snapshot_id, org_id=org_id, project_id=project_id)
+            if _snapshot_done(snapshot, snapshot_id):
+                return snapshot
+            remaining = deadline - time.monotonic() * 1000.0
+            if remaining <= 0:
+                raise NeevAIError(_snapshot_timeout(snapshot, snapshot_id, timeout_ms))
+            time.sleep(min(poll_interval_ms, remaining) / 1000.0)
+
     def delete_snapshot(
         self,
         snapshot_id: str | UUID,
@@ -1067,6 +1119,31 @@ async def get_snapshot(
         raw = await self._client._transport.request("GET", path)
         return coerce_model(Snapshot, raw)
 
+    async def wait_for_snapshot(
+        self,
+        snapshot_id: str | UUID,
+        timeout_ms: int = DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS,
+        poll_interval_ms: int = DEFAULT_SNAPSHOT_POLL_INTERVAL_MS,
+        org_id: str | None = None,
+        project_id: str | None = None,
+    ) -> Snapshot:
+        """Polls a snapshot asynchronously until it is Ready and returns it.
+
+        Raises:
+            NeevAIError: the snapshot Failed, was not Ready within ``timeout_ms``, or a
+                timing is not a positive number.
+        """
+        check_snapshot_wait(timeout_ms, poll_interval_ms)
+        deadline = time.monotonic() * 1000.0 + timeout_ms
+        while True:
+            snapshot = await self.get_snapshot(snapshot_id, org_id=org_id, project_id=project_id)
+            if _snapshot_done(snapshot, snapshot_id):
+                return snapshot
+            remaining = deadline - time.monotonic() * 1000.0
+            if remaining <= 0:
+                raise NeevAIError(_snapshot_timeout(snapshot, snapshot_id, timeout_ms))
+            await asyncio.sleep(min(poll_interval_ms, remaining) / 1000.0)
+
     async def delete_snapshot(
         self,
         snapshot_id: str | UUID,

tests/test_sandboxes.py

@@ -1,5 +1,6 @@
 """Basic sanity tests for the Sandboxes resource (the API, sync)."""
 
+import re
 import uuid
 from unittest.mock import MagicMock
 
@@ -896,3 +897,140 @@ async def test_async_sandboxes_list_bad_label_surfaces_bad_request():
     with pytest.raises(BadRequestError):
         await client.sandboxes.list(labels={"Bad Key": "v"})
     await client.aclose()
+
+
+def _scripted_snapshots(client, statuses: list[dict]) -> list[str]:
+    """Makes get_snapshot return each scripted status in turn, then the last one forever."""
+    calls: list[str] = []
+    models = [Snapshot.model_validate(snapshot_data(**s)) for s in statuses]
+
+    def get_snapshot(snapshot_id, org_id=None, project_id=None):
+        calls.append(str(snapshot_id))
+        return models[min(len(calls) - 1, len(models) - 1)]
+
+    client.sandboxes.get_snapshot = get_snapshot  # type: ignore[method-assign]
+    return calls
+
+
+@pytest.mark.parametrize(
+    ("name", "statuses", "timeout_ms", "expected_error", "expected_polls"),
+    [
+        (
+            "success: ready after polling",
+            [{"status": "Pending"}, {"status": "Running"}, {"status": "Ready"}],
+            5_000,
+            None,
+            3,
+        ),
+        ("success: ready at once", [{"status": "Ready"}], 5_000, None, 1),
+        (
+            "failure: snapshot failed, with its reason",
+            [{"status": "Pending"}, {"status": "Failed", "error_message": "disk full"}],
+            5_000,
+            "failed to capture: disk full",
+            2,
+        ),
+        (
+            "failure: not ready before the timeout",
+            [{"status": "Pending"}],
+            30,
+            "was not Ready within 30ms (status: Pending)",
+            None,
+        ),
+    ],
+)
+def test_sandboxes_wait_for_snapshot(
+    mock_transport, name, statuses, timeout_ms, expected_error, expected_polls
+):
+    client = _make_client(mock_transport)
+    calls = _scripted_snapshots(client, statuses)
+
+    if expected_error is None:
+        snap = client.sandboxes.wait_for_snapshot(
+            "snap-1", timeout_ms=timeout_ms, poll_interval_ms=1
+        )
+        assert snap.status == SnapshotStatus.Ready
+    else:
+        with pytest.raises(NeevAIError, match=re.escape(expected_error)):
+            client.sandboxes.wait_for_snapshot("snap-1", timeout_ms=timeout_ms, poll_interval_ms=1)
+    if expected_polls is not None:
+        assert len(calls) == expected_polls
+    assert set(calls) == {"snap-1"}
+    client.close()
+
+
+@pytest.mark.parametrize(
+    ("timeout_ms", "poll_interval_ms", "bad"),
+    [
+        (0, 1, "timeout_ms"),
+        (-5, 1, "timeout_ms"),
+        (float("nan"), 1, "timeout_ms"),
+        (100, 0, "poll_interval_ms"),
+        (100, float("inf"), "poll_interval_ms"),
+    ],
+)
+def test_sandboxes_wait_for_snapshot_rejects_bad_timings(
+    mock_transport, timeout_ms, poll_interval_ms, bad
+):
+    client = _make_client(mock_transport)
+    calls = _scripted_snapshots(client, [{"status": "Ready"}])
+    with pytest.raises(NeevAIError, match=f"{bad} must be a positive number"):
+        client.sandboxes.wait_for_snapshot(
+            "snap-1", timeout_ms=timeout_ms, poll_interval_ms=poll_interval_ms
+        )
+    assert calls == []
+    client.close()
+
+
+def test_sandbox_handle_snapshot_waits_until_ready(mock_transport):
+    client = _make_client(mock_transport)
+    sb = client.sandboxes.create({"name": "s1", "sandbox_template_id": "sb-ubuntu-26-04-dev"})
+    calls = _scripted_snapshots(client, [{"status": "Running"}, {"status": "Ready"}])
+
+    pending = sb.snapshot({"name": "no-wait"})
+    assert pending.status == SnapshotStatus.Pending
+    assert calls == []
+
+    ready = sb.snapshot({"name": "wait"}, wait_until_ready=True, poll_interval_ms=1)
+    assert ready.status == SnapshotStatus.Ready
+    assert len(calls) == 2
+    client.close()
+
+
+@pytest.mark.asyncio
+async def test_async_sandbox_handle_snapshot_waits_until_ready(async_mock_transport):
+    client = AsyncNeevAI(
+        api_key="test", org_id="org1", project_id="proj1", client=async_mock_transport
+    )
+    sb = await client.sandboxes.create({"name": "s1", "sandbox_template_id": "sb-x"})
+    statuses = [{"status": "Pending"}, {"status": "Failed", "error_message": "no space"}]
+    models = [Snapshot.model_validate(snapshot_data(**s)) for s in statuses]
+    calls: list[str] = []
+
+    async def get_snapshot(snapshot_id, org_id=None, project_id=None):
+        calls.append(str(snapshot_id))
+        return models[min(len(calls) - 1, len(models) - 1)]
+
+    client.sandboxes.get_snapshot = get_snapshot  # type: ignore[method-assign]
+
+    with pytest.raises(NeevAIError, match="failed to capture: no space"):
+        await sb.snapshot({"name": "wait"}, wait_until_ready=True, poll_interval_ms=1)
+    assert len(calls) == 2
+    await client.aclose()
+
+
+def test_sandbox_handle_snapshot_rejects_bad_timing_before_creating(mock_transport):
+    client = _make_client(mock_transport)
+    sb = client.sandboxes.create({"name": "s1", "sandbox_template_id": "sb-ubuntu-26-04-dev"})
+    created: list[str] = []
+
+    def create_snapshot(*args, **kwargs):
+        created.append("called")
+        raise AssertionError("create_snapshot must not run")
+
+    client.sandboxes.create_snapshot = create_snapshot  # type: ignore[method-assign]
+
+    with pytest.raises(NeevAIError, match="timeout_ms must be a positive number"):
+        sb.snapshot({"name": "x"}, wait_until_ready=True, timeout_ms=0)
+    assert created == []
+    client.close()