From 2dccde8797ea3bb1137170f2f377e3f978209b1c Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Mon, 31 Aug 2026 17:27:57 -0700 Subject: [PATCH] feat: Allow setting `ingress_control_config` when creating a `SandboxEnvironmentTemplate` via the SDK (Python, Java, JS). PiperOrigin-RevId: 974166963 --- agentplatform/_genai/sandbox_templates.py | 7 + agentplatform/_genai/types/__init__.py | 6 +- agentplatform/_genai/types/common.py | 191 +++++++++--------- .../unit/agentplatform/genai/test_sandbox.py | 102 +++++++++- 4 files changed, 202 insertions(+), 104 deletions(-) diff --git a/agentplatform/_genai/sandbox_templates.py b/agentplatform/_genai/sandbox_templates.py index d173081f80..66d5767824 100644 --- a/agentplatform/_genai/sandbox_templates.py +++ b/agentplatform/_genai/sandbox_templates.py @@ -62,6 +62,13 @@ def _CreateSandboxEnvironmentTemplateConfig_to_vertex( getv(from_object, ["egress_control_config"]), ) + if getv(from_object, ["ingress_control_config"]) is not None: + setv( + parent_object, + ["ingressControlConfig"], + getv(from_object, ["ingress_control_config"]), + ) + return to_object diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index fc3e0056b5..667ce0b282 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -3457,15 +3457,15 @@ "SandboxEnvironmentTemplateEgressControlConfig", "SandboxEnvironmentTemplateEgressControlConfigDict", "SandboxEnvironmentTemplateEgressControlConfigOrDict", - "CreateSandboxEnvironmentTemplateConfig", - "CreateSandboxEnvironmentTemplateConfigDict", - "CreateSandboxEnvironmentTemplateConfigOrDict", "PSCAutomationConfig", "PSCAutomationConfigDict", "PSCAutomationConfigOrDict", "PrivateServiceConnectConfig", "PrivateServiceConnectConfigDict", "PrivateServiceConnectConfigOrDict", + "CreateSandboxEnvironmentTemplateConfig", + "CreateSandboxEnvironmentTemplateConfigDict", + "CreateSandboxEnvironmentTemplateConfigOrDict", "SandboxEnvironmentTemplate", "SandboxEnvironmentTemplateDict", "SandboxEnvironmentTemplateOrDict", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index ee0f0130c0..0eb2e58739 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -17603,98 +17603,6 @@ class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False): ] -class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for creating a Sandbox Template.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironment - ] = Field( - default=None, - description="""The custom container environment for the sandbox template.""", - ) - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironment - ] = Field( - default=None, - description="""The default container environment for the sandbox template.""", - ) - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( - Field( - default=None, - description="""The egress control config for the sandbox template.""", - ) - ) - - -class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for creating a Sandbox Template.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" - - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironmentDict - ] - """The custom container environment for the sandbox template.""" - - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict - ] - """The default container environment for the sandbox template.""" - - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] - """The egress control config for the sandbox template.""" - - -CreateSandboxEnvironmentTemplateConfigOrDict = Union[ - CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict -] - - -class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): - """Parameters for creating Sandbox Environment Templates.""" - - name: Optional[str] = Field( - default=None, - description="""Name of the agent runtime to create the template under.""", - ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the sandbox template.""" - ) - config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( - default=None, description="""""" - ) - - -class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): - """Parameters for creating Sandbox Environment Templates.""" - - name: Optional[str] - """Name of the agent runtime to create the template under.""" - - display_name: Optional[str] - """The display name of the sandbox template.""" - - config: Optional[CreateSandboxEnvironmentTemplateConfigDict] - """""" - - -_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[ - _CreateSandboxEnvironmentTemplateRequestParameters, - _CreateSandboxEnvironmentTemplateRequestParametersDict, -] - - class PSCAutomationConfig(_common.BaseModel): """PSC config that is used to automatically create PSC endpoints in the user projects.""" @@ -17798,6 +17706,105 @@ class PrivateServiceConnectConfigDict(TypedDict, total=False): ] +class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for creating a Sandbox Template.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironment + ] = Field( + default=None, + description="""The custom container environment for the sandbox template.""", + ) + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironment + ] = Field( + default=None, + description="""The default container environment for the sandbox template.""", + ) + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( + Field( + default=None, + description="""The egress control config for the sandbox template.""", + ) + ) + ingress_control_config: Optional[PrivateServiceConnectConfig] = Field( + default=None, + description="""The ingress control config for the sandbox template.""", + ) + + +class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for creating a Sandbox Template.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironmentDict + ] + """The custom container environment for the sandbox template.""" + + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict + ] + """The default container environment for the sandbox template.""" + + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] + """The egress control config for the sandbox template.""" + + ingress_control_config: Optional[PrivateServiceConnectConfigDict] + """The ingress control config for the sandbox template.""" + + +CreateSandboxEnvironmentTemplateConfigOrDict = Union[ + CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict +] + + +class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for creating Sandbox Environment Templates.""" + + name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime to create the template under.""", + ) + display_name: Optional[str] = Field( + default=None, description="""The display name of the sandbox template.""" + ) + config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( + default=None, description="""""" + ) + + +class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for creating Sandbox Environment Templates.""" + + name: Optional[str] + """Name of the agent runtime to create the template under.""" + + display_name: Optional[str] + """The display name of the sandbox template.""" + + config: Optional[CreateSandboxEnvironmentTemplateConfigDict] + """""" + + +_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _CreateSandboxEnvironmentTemplateRequestParameters, + _CreateSandboxEnvironmentTemplateRequestParametersDict, +] + + class SandboxEnvironmentTemplate(_common.BaseModel): """A sandbox environment template.""" diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index 23eba767db..d128ff5311 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -130,12 +130,10 @@ def test_generate_browser_ws_headers( mock_http_client.request.return_value = genai_types.HttpResponse( body=b'{"endpoint": "test/endpoint"}', headers={} ) - ws_url, headers = ( - self.client.sandboxes.generate_browser_ws_headers( - sandbox_environment=mock_sandbox, - service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL, - timeout=3600, - ) + ws_url, headers = self.client.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.example.vertexai.goog/test/endpoint" assert ( @@ -297,6 +295,94 @@ def test_create_without_spec_template_or_snapshot_raises(self, mock_create): mock_create.assert_not_called() +@pytest.mark.usefixtures("google_auth_mock") +class TestSandboxTemplates: + """Tests for sandbox_templates.SandboxTemplates. + + Exercises the create-side wiring of fields on + `CreateSandboxEnvironmentTemplateConfig`, including + `ingress_control_config` (see cl/972355621), which is otherwise not + covered by the sandboxes.create tests above. + """ + + def setup_method(self): + importlib.reload(initializer) + importlib.reload(aiplatform) + importlib.reload(agentplatform) + self.client = agentplatform.Client( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_create_config_ingress_control_config_is_optional(self): + """Constructing without ingress_control_config must succeed.""" + config = agentplatform_types.CreateSandboxEnvironmentTemplateConfig() + assert config.ingress_control_config is None + + def test_create_config_accepts_ingress_control_config_object(self): + """The typed field accepts a PrivateServiceConnectConfig instance.""" + psc = agentplatform_types.PrivateServiceConnectConfig( + enable_private_service_connect=True, + project_allowlist=["test-project"], + ) + config = agentplatform_types.CreateSandboxEnvironmentTemplateConfig( + ingress_control_config=psc, + ) + assert config.ingress_control_config is not None + assert config.ingress_control_config.enable_private_service_connect is True + assert config.ingress_control_config.project_allowlist == ["test-project"] + + def test_create_config_accepts_ingress_control_config_dict(self): + """The typed field validates and coerces a dict input.""" + config = ( + agentplatform_types.CreateSandboxEnvironmentTemplateConfig.model_validate( + { + "ingress_control_config": { + "enable_private_service_connect": True, + "project_allowlist": ["test-project"], + }, + } + ) + ) + assert isinstance( + config.ingress_control_config, + agentplatform_types.PrivateServiceConnectConfig, + ) + assert config.ingress_control_config.enable_private_service_connect is True + assert config.ingress_control_config.project_allowlist == ["test-project"] + + @mock.patch.object(sandbox_templates.SandboxTemplates, "_create") + def test_create_forwards_ingress_control_config(self, mock_create): + """templates.create(...) passes ingress_control_config through to _create.""" + mock_create.return_value = mock.Mock() + + config = agentplatform_types.CreateSandboxEnvironmentTemplateConfig( + ingress_control_config=agentplatform_types.PrivateServiceConnectConfig( + enable_private_service_connect=True, + project_allowlist=["test-project"], + ), + wait_for_completion=False, + ) + self.client.sandboxes.templates.create( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + display_name="test-template", + config=config, + ) + + mock_create.assert_called_once() + _, kwargs = mock_create.call_args + assert kwargs["name"] == _TEST_AGENT_ENGINE_RESOURCE_NAME + assert kwargs["display_name"] == "test-template" + ingress = kwargs["config"].ingress_control_config + assert ingress is not None + assert ingress.enable_private_service_connect is True + assert ingress.project_allowlist == ["test-project"] + + _MODULES = pytest.mark.parametrize( "module", [sandboxes, vertexai_sandboxes], @@ -412,9 +498,7 @@ def test_generate_access_token_retries_transient_failures(module, status_code): [transient, transient, _ok_response()], ) - with default_patch, session_patch, mock.patch.object( - module.time, "sleep" - ) as sleep: + with default_patch, session_patch, mock.patch.object(module.time, "sleep") as sleep: client_obj = module.Sandboxes(api_client_=mock.Mock()) token = client_obj.generate_access_token( service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL