feat: Add support for wait for operation complete for sandbox sub-modules

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

@@ -21,7 +21,7 @@
 
 
 def test_create_sandbox_snapshot(client):
-    snapshot = client.agent_engines.sandboxes.snapshots.create(
+    snapshot = client.agent_engines.sandboxes.snapshots._create(
         source_sandbox_environment_name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironments/525190525100228608",
         config={
             "display_name": "test_snapshot",
@@ -36,5 +36,5 @@ def test_create_sandbox_snapshot(client):
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
-    test_method="agent_engines.sandboxes.snapshots.create",
+    test_method="agent_engines.sandboxes.snapshots._create",
 )

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

@@ -21,7 +21,7 @@
 
 
 def test_delete_sandbox_snapshot(client):
-    result = client.agent_engines.sandboxes.snapshots.delete(
+    result = client.agent_engines.sandboxes.snapshots._delete(
         name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/421086565159141376",
     )
 
@@ -31,5 +31,5 @@ def test_delete_sandbox_snapshot(client):
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
-    test_method="agent_engines.sandboxes.snapshots.delete",
+    test_method="agent_engines.sandboxes.snapshots._delete",
 )

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

@@ -21,8 +21,8 @@
 
 
 def test_get_sandbox_snapshot(client):
-    snapshot_name = "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/421086565159141376"
-    fetched_snapshot = client.agent_engines.sandboxes.snapshots.get(
+    snapshot_name = "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/2433069698686910464"
+    fetched_snapshot = client.agent_engines.sandboxes.snapshots._get(
         name=snapshot_name,
     )
 
@@ -33,5 +33,5 @@ def test_get_sandbox_snapshot(client):
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
-    test_method="agent_engines.sandboxes.snapshots.get",
+    test_method="agent_engines.sandboxes.snapshots._get",
 )

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

@@ -21,7 +21,7 @@
 
 
 def test_list_sandbox_snapshots(client):
-    snapshots_list_operation = client.agent_engines.sandboxes.snapshots.list(
+    snapshots_list_operation = client.agent_engines.sandboxes.snapshots._list(
         name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472",
     )
 
@@ -37,5 +37,5 @@ def test_list_sandbox_snapshots(client):
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
-    test_method="agent_engines.sandboxes.snapshots.list",
+    test_method="agent_engines.sandboxes.snapshots._list",
 )

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

@@ -50,7 +50,7 @@ def test_sandbox_templates_byoc_create(client):
     }
     sandbox_template_operation = client.agent_engines.sandboxes.templates.create(
         name=(
-            "projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560"
+            "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472"
         ),
         display_name="Test Sandbox Template 1",
         config=config,

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

@@ -27,9 +27,9 @@ def test_sandbox_templates_default_create(client):
             "internet_access": True,
         },
     }
-    sandbox_template_operation = client.agent_engines.sandboxes.templates.create(
+    sandbox_template_operation = client.agent_engines.sandboxes.templates._create(
         name=(
-            "projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560"
+            "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472"
         ),
         display_name="Test Sandbox Template 1",
         config=config,
@@ -42,5 +42,5 @@ def test_sandbox_templates_default_create(client):
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
-    test_method="agent_engines.sandboxes.templates.create",
+    test_method="agent_engines.sandboxes.templates._create",
 )

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

@@ -19,7 +19,7 @@
 
 
 def test_sandbox_templates_list(client):
-    sandbox_templates_list_operation = client.agent_engines.sandboxes.templates.list(
+    sandbox_templates_list_operation = client.agent_engines.sandboxes.templates._list(
         name=(
             "projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560"
         ),
@@ -36,5 +36,5 @@ def test_sandbox_templates_list(client):
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
-    test_method="agent_engines.sandboxes.templates.list",
+    test_method="agent_engines.sandboxes.templates._list",
 )

vertexai/_genai/sandbox_snapshots.py

@@ -15,20 +15,25 @@
 
 # Code generated by the Google Gen AI SDK generator DO NOT EDIT.
 
+import functools
 import json
 import logging
-from typing import Any, Optional, Union
+from typing import Any, Iterator, Optional, Union
 from urllib.parse import urlencode
 
 from google.genai import _api_module
 from google.genai import _common
 from google.genai._common import get_value_by_path as getv
 from google.genai._common import set_value_by_path as setv
+from google.genai.pagers import Pager
 
+from . import _agent_engines_utils
 from . import types
 
 logger = logging.getLogger("vertexai_genai.sandboxsnapshots")
 
+logger.setLevel(logging.INFO)
+
 
 def _CreateAgentEngineSandboxSnapshotConfig_to_vertex(
     from_object: Union[dict[str, Any], object],
@@ -144,14 +149,21 @@ def _ListSandboxEnvironmentSnapshotsRequestParameters_to_vertex(
 class SandboxSnapshots(_api_module.BaseModule):
     """Sandbox environment snapshot commands."""
 
-    def create(
+    def _create(
         self,
         *,
         source_sandbox_environment_name: str,
         config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None,
     ) -> types.AgentEngineSandboxSnapshotOperation:
         """
-        Snapshots a sandbox environment.
+        Snapshots an existing sandbox environment.
+
+        Args:
+            source_sandbox_environment_name (str):
+                Required. The name of the sandbox environment to snapshot.
+                projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id}
+            config (CreateAgentEngineSandboxSnapshotConfig):
+                Optional. The configuration for the sandbox snapshot.
 
         """
 
@@ -219,7 +231,7 @@ def create(
         self._api_client._verify_response(return_value)
         return return_value
 
-    def delete(
+    def _delete(
         self,
         *,
         name: str,
@@ -294,7 +306,7 @@ def delete(
         self._api_client._verify_response(return_value)
         return return_value
 
-    def get(
+    def _get(
         self,
         *,
         name: str,
@@ -369,7 +381,7 @@ def get(
         self._api_client._verify_response(return_value)
         return return_value
 
-    def list(
+    def _list(
         self,
         *,
         name: str,
@@ -444,7 +456,7 @@ def list(
         self._api_client._verify_response(return_value)
         return return_value
 
-    def _get_sandbox_snapshot_operation(
+    def get_sandbox_snapshot_operation(
         self,
         *,
         operation_name: str,
@@ -514,18 +526,126 @@ def _get_sandbox_snapshot_operation(
         self._api_client._verify_response(return_value)
         return return_value
 
+    def create(
+        self,
+        *,
+        source_sandbox_environment_name: str,
+        config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None,
+        poll_interval_seconds: float = 0.1,
+    ) -> types.AgentEngineSandboxSnapshotOperation:
+        """Snapshots an existing sandbox environment.
+
+        Args:
+            source_sandbox_environment_name (str):
+                Required. The name of the sandbox environment to snapshot.
+                projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id}
+            config (CreateAgentEngineSandboxSnapshotConfig):
+                Optional. The configuration for the sandbox snapshot.
+            poll_interval_seconds (int):
+                Optional. Seconds to wait between polling for operation status. Defaults to 0.1.
+
+        Returns:
+            AgentEngineSandboxSnapshotOperation: The operation for creating the sandbox snapshot.
+        """
+        operation = self._create(
+            source_sandbox_environment_name=source_sandbox_environment_name,
+            config=config,
+        )
+        if config is None:
+            config = types.CreateAgentEngineSandboxSnapshotConfig()
+        elif isinstance(config, dict):
+            config = types.CreateAgentEngineSandboxSnapshotConfig.model_validate(config)
+        if config.wait_for_completion:
+            if not operation.done:
+                operation = _agent_engines_utils._await_operation(
+                    operation_name=operation.name,
+                    get_operation_fn=self.get_sandbox_snapshot_operation,
+                    poll_interval_seconds=poll_interval_seconds,
+                )
+            # We need to make a call to get the sandbox snapshot because the operation
+            # response might not contain the relevant fields.
+            if not operation.response:
+                raise ValueError("Error retrieving sandbox snapshot.")
+            operation.response = self.get(name=operation.response.name)
+        return operation
+
+    def list(
+        self,
+        *,
+        name: str,
+        config: Optional[types.ListSandboxEnvironmentSnapshotsConfigOrDict] = None,
+    ) -> Iterator[types.SandboxEnvironmentSnapshot]:
+        """Lists Agent Engine sandbox snapshots.
+
+        Args:
+            name (str):
+                Required. The name of the agent engine to list sandbox snapshots for.
+                projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{sandbox_snapshot_id}
+            config (ListSandboxEnvironmentSnapshotsConfig):
+                Optional. The configuration for the sandbox snapshots to list.
+
+        Returns:
+            Iterable[SandboxEnvironmentSnapshot]: An iterable of agent engine sandbox snapshots.
+        """
+        return Pager(
+            "sandbox_environment_snapshots",
+            functools.partial(self._list, name=name),
+            self._list(name=name, config=config),
+            config,
+        )
+
+    def get(
+        self,
+        *,
+        name: str,
+        config: Optional[types.GetSandboxEnvironmentSnapshotConfigOrDict] = None,
+    ) -> types.SandboxEnvironmentSnapshot:
+        """Gets a sandbox snapshot in the Agent Engine.
+        Args:
+          name (str):
+              Required. A fully-qualified resource name or ID such as
+              projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{snapshot_id}
+              or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{snapshot_id}".
+          config (GetSandboxEnvironmentSnapshotConfigOrDict):
+              Optional. The configuration for the sandbox snapshot to get.
+        """
+        return self._get(name=name, config=config)
+
+    def delete(
+        self,
+        *,
+        name: str,
+        config: Optional[types.DeleteSandboxEnvironmentSnapshotConfigOrDict] = None,
+    ) -> types.DeleteSandboxEnvironmentSnapshotOperation:
+        """Deletes a sandbox snapshot in the Agent Engine.
+        Args:
+            name (str):
+                Required. The name of the sandbox snapshot to delete.
+                Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{snapshot_id}
+            config (DeleteSandboxEnvironmentSnapshotConfigOrDict):
+                Optional. Configuration for the delete operation.
+        """
+        return self._delete(name=name, config=config)
+
 
 class AsyncSandboxSnapshots(_api_module.BaseModule):
     """Sandbox environment snapshot commands."""
 
-    async def create(
+    async def _create(
         self,
         *,
         source_sandbox_environment_name: str,
         config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None,
     ) -> types.AgentEngineSandboxSnapshotOperation:
         """
-        Snapshots a sandbox environment.
+        Snapshots an existing sandbox environment.
+
+        Args:
+            source_sandbox_environment_name (str):
+                Required. The name of the sandbox environment to snapshot.
+                projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id}
+            config (CreateAgentEngineSandboxSnapshotConfig):
+                Optional. The configuration for the sandbox snapshot.
 
         """
 
@@ -595,7 +715,7 @@ async def create(
         self._api_client._verify_response(return_value)
         return return_value
 
-    async def delete(
+    async def _delete(
         self,
         *,
         name: str,
@@ -672,7 +792,7 @@ async def delete(
         self._api_client._verify_response(return_value)
         return return_value
 
-    async def get(
+    async def _get(
         self,
         *,
         name: str,
@@ -749,7 +869,7 @@ async def get(
         self._api_client._verify_response(return_value)
         return return_value
 
-    async def list(
+    async def _list(
         self,
         *,
         name: str,
@@ -826,7 +946,7 @@ async def list(
         self._api_client._verify_response(return_value)
         return return_value
 
-    async def _get_sandbox_snapshot_operation(
+    async def get_sandbox_snapshot_operation(
         self,
         *,
         operation_name: str,