From 07f57068f28350af24831abf87c83ae0ccd0a96c Mon Sep 17 00:00:00 2001
From: A Vertex SDK engineer <vertex-sdk-bot@google.com>
Date: Mon, 27 Apr 2026 21:17:57 -0700
Subject: [PATCH] feat: faster sandbox creation with templates and snapshots
 and improve dataplane routing and security.

Multitenancy Sandbox support

PiperOrigin-RevId: 906714249
---
 tests/unit/vertexai/genai/test_sandbox.py | 44 +++++++++++++++++++++--
 vertexai/_genai/types/common.py           | 25 +++++++++++++
 2 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/tests/unit/vertexai/genai/test_sandbox.py b/tests/unit/vertexai/genai/test_sandbox.py
index cc8cd7968b..45949ac60d 100644
--- a/tests/unit/vertexai/genai/test_sandbox.py
+++ b/tests/unit/vertexai/genai/test_sandbox.py
@@ -73,8 +73,9 @@ def teardown_method(self):
     @mock.patch.object(client.Client, "_get_api_client")
     def test_send_command(self, mock_get_api_client):
         mock_sandbox = mock.Mock()
-        mock_sandbox.connection_info.load_balancer_ip = "127.0.0.1"
-        mock_sandbox.connection_info.load_balancer_hostname = None
+        mock_sandbox.connection_info.load_balancer_ip = None
+        mock_sandbox.connection_info.load_balancer_hostname = "test-us-central1.autopush-sandbox.vertexai.goog"
+        mock_sandbox.connection_info.routing_token = "test_routing_token"
         mock_http_client = mock_get_api_client.return_value
         mock_http_client.request.return_value = genai_types.HttpResponse(
             body=b"{}", headers={}
@@ -91,7 +92,44 @@ def test_send_command(self, mock_get_api_client):
         assert call_args is not None
         _, kwargs = call_args
         http_options = kwargs["http_options"]
-        assert http_options.base_url == "http://127.0.0.1/test/path"
+        assert http_options.base_url == (
+            "https://test-us-central1.autopush-sandbox.vertexai.goog/test/path"
+        )
         assert http_options.headers["Authorization"] == "Bearer test_token"
 
         mock_http_client.request.assert_called_with("GET", "test/path", {})
+
+    @mock.patch(
+        "google.cloud.aiplatform.vertexai._genai.sandboxes.Sandboxes.generate_access_token"
+    )
+    @mock.patch.object(client.Client, "_get_api_client")
+    def test_generate_browser_ws_headers(
+        self, mock_get_api_client, mock_generate_access_token
+    ):
+        mock_generate_access_token.return_value = "test_token"
+
+        mock_sandbox = mock.Mock()
+        mock_sandbox.connection_info.load_balancer_ip = None
+        mock_sandbox.connection_info.load_balancer_hostname = (
+            "test-us-central1.autopush-sandbox.vertexai.goog"
+        )
+        mock_sandbox.connection_info.routing_token = "test_routing_token"
+        mock_http_client = mock_get_api_client.return_value
+        mock_http_client.request.return_value = genai_types.HttpResponse(
+            body=b'{"endpoint": "test/endpoint"}', headers={}
+        )
+        ws_url, headers = (
+            self.client.agent_engines.sandboxes.generate_browser_ws_headers(
+                sandbox_environment=mock_sandbox,
+                service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL,
+                timeout=3600,
+            )
+        )
+        assert (
+            ws_url
+            == "wss://test-us-central1.autopush-sandbox.vertexai.goog/test/endpoint"
+        )
+        assert (
+            headers["Sec-WebSocket-Protocol"]
+            == "v1.stream, test_token, test_routing_token, 9222"
+        )
diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py
index b2b26b8573..cf7b3f94cf 100644
--- a/vertexai/_genai/types/common.py
+++ b/vertexai/_genai/types/common.py
@@ -11733,6 +11733,20 @@ class CreateAgentEngineSandboxConfig(_common.BaseModel):
         default=None,
         description="""The TTL for this resource. The expiration time is computed: now + TTL.""",
     )
+    sandbox_environment_template: Optional[str] = Field(
+        default=None,
+        description="""The name of the sandbox environment template to create the sandbox from. The sandbox environment template should be in the format:
+      projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}""",
+    )
+    sandbox_environment_snapshot: Optional[str] = Field(
+        default=None,
+        description="""The name of the sandbox environment snapshot to restore the sandbox from. The sandbox environment snapshot should be in the format:
+      projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}""",
+    )
+    owner: Optional[str] = Field(
+        default=None,
+        description="""Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""",
+    )
 
 
 class CreateAgentEngineSandboxConfigDict(TypedDict, total=False):
@@ -11753,6 +11767,17 @@ class CreateAgentEngineSandboxConfigDict(TypedDict, total=False):
     ttl: Optional[str]
     """The TTL for this resource. The expiration time is computed: now + TTL."""
 
+    sandbox_environment_template: Optional[str]
+    """The name of the sandbox environment template to create the sandbox from. The sandbox environment template should be in the format:
+      projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}"""
+
+    sandbox_environment_snapshot: Optional[str]
+    """The name of the sandbox environment snapshot to restore the sandbox from. The sandbox environment snapshot should be in the format:
+      projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}"""
+
+    owner: Optional[str]
+    """Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner."""
+
 
 CreateAgentEngineSandboxConfigOrDict = Union[
     CreateAgentEngineSandboxConfig, CreateAgentEngineSandboxConfigDict