feat: GenAI SDK client - Support agent engine sandbox http request in genai sdk

PiperOrigin-RevId: 816865842

Author: vertex-sdk-bot

setup.py

@@ -168,6 +168,7 @@
     "opentelemetry-exporter-otlp-proto-http < 2",
     "pydantic >= 2.11.1, < 3",
     "typing_extensions",
+    "google-cloud-iam",
 ]
 
 evaluation_extra_require = [

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

@@ -0,0 +1,62 @@
+# Copyright 2025 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
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+
+
+def test_send_command_sandbox(client):
+    # agent_engine = client.agent_engines.create()
+    # assert isinstance(agent_engine, types.AgentEngine)
+    # assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+    # agent_engine_name = "projects/254005681254/locations/us-central1/reasoningEngines/2112984271655272448"
+    # operation = client.agent_engines.sandboxes.create(
+    #     # name=agent_engine.api_resource.name,
+    #     name=agent_engine_name,
+    #     spec={
+    #         "code_execution_environment": {
+    #             "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB"
+    #         }
+    #     },
+    #     config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"),
+    # )
+    # assert isinstance(operation, types.AgentEngineSandboxOperation)
+
+    client._api_client.project = None
+    client._api_client.location = None
+    client._api_client.api_version = None
+    token = client.agent_engines.sandboxes.generate_access_token(
+        service_account_email="sign-verify-jwt@mariner-proxy.iam.gserviceaccount.com",
+        sandbox_id="tenghuil-manual-sandbox-new",
+    )
+    response = client.agent_engines.sandboxes.send_command(
+        http_method="GET",
+        path="/",
+        query_params={},
+        access_token=token,
+        headers={},
+        request_dict={},
+        sandbox_environment=None,
+    )
+    # assert token == "xxx"
+    assert response.body == "Hello World"
+    # assert response.outputs[0].mime_type == "application/json"
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.sandboxes.send_command",
+)

vertexai/_genai/sandboxes.py

@@ -19,11 +19,16 @@
 import json
 import logging
 import mimetypes
+import secrets
+import time
 from typing import Any, Iterator, Optional, Union
 from urllib.parse import urlencode
 
+from google import genai
+from google.cloud import iam_credentials_v1
 from google.genai import _api_module
 from google.genai import _common
+from google.genai import types as genai_types
 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
@@ -704,6 +709,111 @@ def delete(
         """
         return self._delete(name=name, config=config)
 
+    def generate_access_token(
+        self,
+        service_account_email: str,
+        sandbox_id: str,
+        port: str = "8080",
+        timeout: int = 3600,
+    ) -> str:
+        """Signs a JWT with a Google Cloud service account."""
+        client = iam_credentials_v1.IAMCredentialsClient()
+        name = f"projects/-/serviceAccounts/{service_account_email}"
+        custom_claims = {"port": port, "sandbox_id": sandbox_id}
+        payload = {
+            "iat": int(time.time()),
+            "exp": int(time.time()) + timeout,
+            "iss": service_account_email,
+            "nonce": secrets.randbelow(1000000000) + 1,
+            "aud": "vmaas-proxy-api",  # default audience for sandbox proxy
+            **custom_claims,
+        }
+        request = iam_credentials_v1.SignJwtRequest(
+            name=name,
+            payload=json.dumps(payload),
+        )
+        response = client.sign_jwt(request=request)
+        return response.signed_jwt
+
+    def send_command(
+        self,
+        http_method: str,
+        access_token: str,
+        sandbox_environment: types.SandboxEnvironment,
+        path: str = None,
+        query_params: Optional[dict[str, object]] = None,
+        headers: Optional[dict[str, str]] = None,
+        request_dict: Optional[dict[str, object]] = None,
+    ) -> genai_types.HttpResponse:
+        """Sends a command to the sandbox."""
+        headers = headers or {}
+        request_dict = request_dict or {}
+        connection_info = sandbox_environment.connection_info
+        if not connection_info:
+            raise ValueError("Connection info is not available.")
+        if connection_info.load_balancer_hostname:
+            endpoint = "https://" + connection_info.load_balancer_hostname
+        elif connection_info.load_balancer_ip:
+            endpoint = "http://" + connection_info.load_balancer_ip
+        else:
+            raise ValueError("Load balancer hostname or ip is not available.")
+
+        path = path or ""
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        headers["Authorization"] = f"Bearer {access_token}"
+        endpoint = endpoint + path if path.startswith("/") else endpoint + "/" + path
+        http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint)
+        http_client = genai.Client(vertexai=True, http_options=http_options)
+        response = http_client._api_client.request(http_method, path, request_dict)
+        return genai_types.HttpResponse(
+            headers=response.headers,
+            body=response.body,
+        )
+
+    def generate_browser_ws_headers(
+        self,
+        sandbox_environment: types.SandboxEnvironment,
+        service_account_email: str,
+        timeout: int = 3600,
+    ) -> tuple[str, dict[str, str]]:
+        """Generates the websocket upgrade headers for the browser."""
+        sandbox_id = sandbox_environment.name
+        # port 8080 is the default port for http endpoint.
+        http_access_token = self.generate_access_token(
+            service_account_email, sandbox_id, "8080", timeout
+        )
+        response = self.send_command(
+            "GET",
+            http_access_token,
+            sandbox_environment,
+            path="/cdp_ws_endpoint",
+        )
+        if not response:
+            raise ValueError("Failed to get the websocket endpoint.")
+        body_dict = json.loads(response.body)
+        ws_path = body_dict["endpoint"]
+
+        ws_url = "wss://test-us-central1.autopush-sandbox.vertexai.goog"
+        if sandbox_environment and sandbox_environment.connection_info:
+            connection_info = sandbox_environment.connection_info
+            if connection_info.load_balancer_hostname:
+                ws_url = "wss://" + connection_info.load_balancer_hostname
+            elif connection_info.load_balancer_ip:
+                ws_url = "ws://" + connection_info.load_balancer_ip
+            else:
+                raise ValueError("Load balancer hostname or ip is not available.")
+        ws_url = ws_url + "/" + ws_path
+
+        # port 9222 is the default port for the browser websocket endpoint.
+        ws_access_token = self.generate_access_token(
+            service_account_email, sandbox_id, "9222", timeout
+        )
+
+        headers = {}
+        headers["Sec-WebSocket-Protocol"] = f"binary, {ws_access_token}"
+        return ws_url, headers
+
 
 class AsyncSandboxes(_api_module.BaseModule):