feat: Add support for snapshots for Sandboxes in Vertex AI GenAI SDK.

PiperOrigin-RevId: 897294488

Author: vertex-sdk-bot

tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_create.py

@@ -0,0 +1,55 @@
+"""Tests for Sandbox Snapshot create operation."""
+
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring,bad-indentation
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_create_sandbox_snapshot(client):
+  # 1. Create an Agent Engine
+  agent_engine = client.agent_engines.create()
+
+  # 2. Create a Sandbox
+  sandbox_op = client.agent_engines.sandboxes.create(
+      name=agent_engine.api_resource.name,
+      poll_interval_seconds=1,
+      spec={
+          "code_execution_environment": {
+              "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB"
+          }
+      },
+      config=types.CreateAgentEngineSandboxConfig(
+          display_name="test_sandbox", ttl="3600s"
+      ),
+  )
+  sandbox = sandbox_op.response
+
+  # 3. Create a Snapshot
+  snapshot = client.agent_engines.sandboxes.snapshots.create(
+      name=sandbox.name,
+      sandbox_environment_snapshot={},
+  )
+
+  assert isinstance(snapshot, types.SandboxEnvironmentSnapshot)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.sandboxes.snapshots.create",
+)

tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_delete.py

@@ -0,0 +1,60 @@
+"""Tests for Sandbox Snapshot delete operation."""
+
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring,bad-indentation
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_delete_sandbox_snapshot(client):
+  # 1. Create an Agent Engine
+  agent_engine = client.agent_engines.create()
+
+  # 2. Create a Sandbox
+  sandbox_op = client.agent_engines.sandboxes.create(
+      name=agent_engine.api_resource.name,
+      poll_interval_seconds=1,
+      spec={
+          "code_execution_environment": {
+              "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB"
+          }
+      },
+      config=types.CreateAgentEngineSandboxConfig(
+          display_name="test_sandbox", ttl="3600s"
+      ),
+  )
+  sandbox = sandbox_op.response
+
+  # 3. Create a Snapshot
+  snapshot = client.agent_engines.sandboxes.snapshots.create(
+      name=sandbox.name,
+      sandbox_environment_snapshot={},
+  )
+
+  # 4. Delete the Snapshot
+  result = client.agent_engines.sandboxes.snapshots.delete(
+      name=snapshot.name,
+  )
+
+  assert result is None
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.sandboxes.snapshots.delete",
+)

tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_get.py

@@ -0,0 +1,61 @@
+"""Tests for Sandbox Snapshot get operation."""
+
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring,bad-indentation
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_get_sandbox_snapshot(client):
+  # 1. Create an Agent Engine
+  agent_engine = client.agent_engines.create()
+
+  # 2. Create a Sandbox
+  sandbox_op = client.agent_engines.sandboxes.create(
+      name=agent_engine.api_resource.name,
+      poll_interval_seconds=1,
+      spec={
+          "code_execution_environment": {
+              "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB"
+          }
+      },
+      config=types.CreateAgentEngineSandboxConfig(
+          display_name="test_sandbox", ttl="3600s"
+      ),
+  )
+  sandbox = sandbox_op.response
+
+  # 3. Create a Snapshot
+  snapshot = client.agent_engines.sandboxes.snapshots.create(
+      name=sandbox.name,
+      sandbox_environment_snapshot={},
+  )
+
+  # 4. Get the Snapshot
+  fetched_snapshot = client.agent_engines.sandboxes.snapshots.get(
+      name=snapshot.name,
+  )
+
+  assert isinstance(fetched_snapshot, types.SandboxEnvironmentSnapshot)
+  assert fetched_snapshot.name == snapshot.name
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.sandboxes.snapshots.get",
+)

tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_list.py

@@ -0,0 +1,66 @@
+"""Tests for Sandbox Snapshot list operation."""
+
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring,bad-indentation
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai._genai import types
+
+
+def test_list_sandbox_snapshots(client):
+  # 1. Create an Agent Engine
+  agent_engine = client.agent_engines.create()
+
+  # 2. Create a Sandbox
+  sandbox_op = client.agent_engines.sandboxes.create(
+      name=agent_engine.api_resource.name,
+      poll_interval_seconds=1,
+      spec={
+          "code_execution_environment": {
+              "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB"
+          }
+      },
+      config=types.CreateAgentEngineSandboxConfig(
+          display_name="test_sandbox", ttl="3600s"
+      ),
+  )
+  sandbox = sandbox_op.response
+
+  # 3. Create a Snapshot
+  client.agent_engines.sandboxes.snapshots.create(
+      name=sandbox.name,
+      sandbox_environment_snapshot={},
+  )
+
+  # 4. List Snapshots
+  snapshots_pager = client.agent_engines.sandboxes.snapshots.list(
+      name=sandbox.name,
+  )
+
+  # Verify it returns a pager and we can iterate over it
+  count = 0
+  for snapshot in snapshots_pager:
+    assert isinstance(snapshot, types.SandboxEnvironmentSnapshot)
+    count += 1
+
+  assert count >= 1
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.sandboxes.snapshots.list",
+)