diff --git a/.changeset/snapshot-wait.md b/.changeset/snapshot-wait.md new file mode 100644 index 0000000..f266c55 --- /dev/null +++ b/.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. diff --git a/src/neevai/handles/sandbox.py b/src/neevai/handles/sandbox.py index da58212..381bba9 100644 --- a/src/neevai/handles/sandbox.py +++ b/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]: diff --git a/src/neevai/resources/sandboxes.py b/src/neevai/resources/sandboxes.py index 794c388..dc31afd 100644 --- a/src/neevai/resources/sandboxes.py +++ b/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, diff --git a/tests/test_sandboxes.py b/tests/test_sandboxes.py index 2afd9fa..a106caf 100644 --- a/tests/test_sandboxes.py +++ b/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()