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

PiperOrigin-RevId: 816865842

Author: vertex-sdk-bot

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 random
+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,52 @@ 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": random.randint(1, 1000000000),
+            "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,
+        path: str,
+        query_params: Any,
+        access_token: str,
+        headers: dict[str, str],
+        request_dict: dict[str, object],
+        sandbox_environment: Optional[types.SandboxEnvironment] = None,
+    ) -> str | None:
+        """Sends a command to the sandbox."""
+        # TODO(tenghuil): Get connection info from sandbox environment when it's ready.
+        endpoint = "https://test-us-central1.autopush-sandbox.vertexai.goog/" + path
+
+        headers = {"Authorization": f"Bearer {access_token}"}
+        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 response
+
 
 class AsyncSandboxes(_api_module.BaseModule):