From 4421a02614e7c648e5d97d89b34da10e194b5cc8 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 13 Nov 2026 14:17:56 -0800 Subject: [PATCH] fix: allow clearing `env_vars` / `secret_env` and updating `deployment_spec` without `agent` or source options in `runtimes.update()` and `agent_engines.update()` PiperOrigin-RevId: 1021322798 --- agentplatform/_genai/runtimes.py | 54 ++++--- .../agentplatform/genai/test_agent_engines.py | 152 +++++++++++++++++- vertexai/_genai/agent_engines.py | 54 ++++--- 3 files changed, 215 insertions(+), 45 deletions(-) diff --git a/agentplatform/_genai/runtimes.py b/agentplatform/_genai/runtimes.py index 24f5333d60..0ba49c550e 100644 --- a/agentplatform/_genai/runtimes.py +++ b/agentplatform/_genai/runtimes.py @@ -2277,6 +2277,7 @@ def _create_config( update_masks.append("traffic_config") config["traffic_config"] = traffic_config + caller_provided_env_vars = env_vars is not None if agent_framework == "google-adk": env_vars = _runtimes_utils._add_telemetry_enablement_env(env_vars) @@ -2393,33 +2394,42 @@ def _create_config( or container_concurrency is not None or keep_alive_probe is not None ) - if runtime_spec is None and is_deployment_spec_updated: + has_code_or_container_spec = runtime_spec is not None + if ( + mode == "create" + and not has_code_or_container_spec + and is_deployment_spec_updated + ): raise ValueError( - "To update `env_vars`, `psc_interface_config`, `min_instances`, " + "To set `env_vars`, `psc_interface_config`, `min_instances`, " "`max_instances`, `resource_limits`, `container_concurrency`, or " "`keep_alive_probe`, you must also provide the `agent` variable or " "the source code options (`source_packages`, " "`developer_connect_source`, `git_source` or `agent_config_source`)." ) - if runtime_spec is not None: - if is_deployment_spec_updated: - ( - deployment_spec, - deployment_update_masks, - ) = self._generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - agent_gateway_config=agent_gateway_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - keep_alive_probe=keep_alive_probe, - ) + if is_deployment_spec_updated: + ( + deployment_spec, + deployment_update_masks, + ) = self._generate_deployment_spec_or_raise( + env_vars=env_vars, + caller_provided_env_vars=caller_provided_env_vars, + psc_interface_config=psc_interface_config, + agent_gateway_config=agent_gateway_config, + min_instances=min_instances, + max_instances=max_instances, + resource_limits=resource_limits, + container_concurrency=container_concurrency, + keep_alive_probe=keep_alive_probe, + ) + if deployment_update_masks: + if runtime_spec is None: + runtime_spec = {} update_masks.extend(deployment_update_masks) runtime_spec["deployment_spec"] = deployment_spec + if has_code_or_container_spec: if agent_server_mode: if not runtime_spec.get("deployment_spec"): runtime_spec["deployment_spec"] = ( @@ -2498,6 +2508,7 @@ def _generate_deployment_spec_or_raise( self, *, env_vars: Optional[dict[str, Union[str, Any]]] = None, + caller_provided_env_vars: bool = True, psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, agent_gateway_config: Optional[ types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict @@ -2510,9 +2521,10 @@ def _generate_deployment_spec_or_raise( ) -> Tuple[dict[str, Any], Sequence[str]]: deployment_spec: dict[str, Any] = {} update_masks = [] - if env_vars: + if env_vars is not None: deployment_spec["env"] = [] - deployment_spec["secret_env"] = [] + if caller_provided_env_vars: + deployment_spec["secret_env"] = [] if isinstance(env_vars, dict): self._update_deployment_spec_with_env_vars_dict_or_raise( deployment_spec=deployment_spec, @@ -2520,9 +2532,9 @@ def _generate_deployment_spec_or_raise( ) else: raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.") - if deployment_spec.get("env"): + if caller_provided_env_vars or deployment_spec.get("env"): update_masks.append("spec.deployment_spec.env") - if deployment_spec.get("secret_env"): + if caller_provided_env_vars or deployment_spec.get("secret_env"): update_masks.append("spec.deployment_spec.secret_env") if psc_interface_config: deployment_spec["psc_interface_config"] = psc_interface_config diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index aa5c5114aa..afb16e0e3a 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -3252,16 +3252,162 @@ def test_update_runtime_extra_packages( @mock.patch.object(_runtimes_utils, "_prepare") @mock.patch.object(_runtimes_utils, "_await_operation") - def test_update_runtime_deployment_config_without_agent_raises( + def test_create_runtime_deployment_config_without_agent_raises( self, mock_await_operation, mock_prepare ): - with pytest.raises(ValueError, match="To update `env_vars`"): + with pytest.raises(ValueError, match="To set `env_vars`"): + self.client.runtimes.create( + config=_genai_types.AgentRuntimeConfig( + env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT + ), + ) + + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_deployment_config_without_agent( + self, mock_await_operation + ): + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config=_genai_types.AgentRuntimeConfig( - env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT + env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, ), ) + update_mask = ",".join( + [ + "spec.deployment_spec.env", + "spec.deployment_spec.secret_env", + ] + ) + query_params = {"updateMask": update_mask} + request_mock.assert_called_with( + "patch", + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}?{urlencode(query_params)}", + { + "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, + "spec": { + "deployment_spec": _TEST_AGENT_ENGINE_SPEC.get( + "deployment_spec" + ), + }, + "_query": {"updateMask": update_mask}, + }, + None, + ) + + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_clear_secret_env(self, mock_await_operation): + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.runtimes.update( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + config=_genai_types.AgentRuntimeConfig( + env_vars={"TEST_ENV_VAR": "TEST_ENV_VAR_VALUE"}, + ), + ) + update_mask = ",".join( + [ + "spec.deployment_spec.env", + "spec.deployment_spec.secret_env", + ] + ) + query_params = {"updateMask": update_mask} + request_mock.assert_called_with( + "patch", + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}?{urlencode(query_params)}", + { + "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, + "spec": { + "deployment_spec": { + "env": [ + { + "name": "TEST_ENV_VAR", + "value": "TEST_ENV_VAR_VALUE", + } + ], + "secret_env": [], + }, + }, + "_query": {"updateMask": update_mask}, + }, + None, + ) + + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_clear_all_env_vars(self, mock_await_operation): + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.runtimes.update( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + config=_genai_types.AgentRuntimeConfig(env_vars={}), + ) + update_mask = ",".join( + [ + "spec.deployment_spec.env", + "spec.deployment_spec.secret_env", + ] + ) + query_params = {"updateMask": update_mask} + request_mock.assert_called_with( + "patch", + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}?{urlencode(query_params)}", + { + "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME}, + "spec": { + "deployment_spec": { + "env": [], + "secret_env": [], + }, + }, + "_query": {"updateMask": update_mask}, + }, + None, + ) + + def test_update_runtime_adk_framework_without_env_vars_preserves_env(self): + config = self.client.runtimes._create_config( + mode="update", + agent_framework="google-adk", + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + ) + assert "secret_env" not in config["spec"]["deployment_spec"] + assert "spec.deployment_spec.secret_env" not in config["update_mask"] + assert "spec.deployment_spec.env" in config["update_mask"] + + def test_update_runtime_empty_deployment_field_without_mask_omits_spec( + self, + ): + config = self.client.runtimes._create_config( + mode="update", + psc_interface_config={}, + ) + assert "spec" not in config + assert "update_mask" not in config @mock.patch.object(_runtimes_utils, "_prepare") @mock.patch.object(_runtimes_utils, "_await_operation") diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 896e3ca95e..1cc970d802 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -2358,6 +2358,7 @@ def _create_config( update_masks.append("traffic_config") config["traffic_config"] = traffic_config + caller_provided_env_vars = env_vars is not None if agent_framework == "google-adk": env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) @@ -2462,33 +2463,42 @@ def _create_config( or container_concurrency is not None or keep_alive_probe is not None ) - if agent_engine_spec is None and is_deployment_spec_updated: + has_code_or_container_spec = agent_engine_spec is not None + if ( + mode == "create" + and not has_code_or_container_spec + and is_deployment_spec_updated + ): raise ValueError( - "To update `env_vars`, `psc_interface_config`, `min_instances`, " + "To set `env_vars`, `psc_interface_config`, `min_instances`, " "`max_instances`, `resource_limits`, `container_concurrency`, or " "`keep_alive_probe`, you must also provide the `agent` variable or " "the source code options (`source_packages`, " "`developer_connect_source` or `agent_config_source`)." ) - if agent_engine_spec is not None: - if is_deployment_spec_updated: - ( - deployment_spec, - deployment_update_masks, - ) = self._generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - agent_gateway_config=agent_gateway_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - keep_alive_probe=keep_alive_probe, - ) + if is_deployment_spec_updated: + ( + deployment_spec, + deployment_update_masks, + ) = self._generate_deployment_spec_or_raise( + env_vars=env_vars, + caller_provided_env_vars=caller_provided_env_vars, + psc_interface_config=psc_interface_config, + agent_gateway_config=agent_gateway_config, + min_instances=min_instances, + max_instances=max_instances, + resource_limits=resource_limits, + container_concurrency=container_concurrency, + keep_alive_probe=keep_alive_probe, + ) + if deployment_update_masks: + if agent_engine_spec is None: + agent_engine_spec = {} update_masks.extend(deployment_update_masks) agent_engine_spec["deployment_spec"] = deployment_spec + if has_code_or_container_spec: if agent_server_mode: if not agent_engine_spec.get("deployment_spec"): agent_engine_spec["deployment_spec"] = ( @@ -2564,6 +2574,7 @@ def _generate_deployment_spec_or_raise( self, *, env_vars: Optional[dict[str, Union[str, Any]]] = None, + caller_provided_env_vars: bool = True, psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, agent_gateway_config: Optional[ types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict @@ -2576,9 +2587,10 @@ def _generate_deployment_spec_or_raise( ) -> Tuple[dict[str, Any], Sequence[str]]: deployment_spec: dict[str, Any] = {} update_masks = [] - if env_vars: + if env_vars is not None: deployment_spec["env"] = [] - deployment_spec["secret_env"] = [] + if caller_provided_env_vars: + deployment_spec["secret_env"] = [] if isinstance(env_vars, dict): self._update_deployment_spec_with_env_vars_dict_or_raise( deployment_spec=deployment_spec, @@ -2586,9 +2598,9 @@ def _generate_deployment_spec_or_raise( ) else: raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.") - if deployment_spec.get("env"): + if caller_provided_env_vars or deployment_spec.get("env"): update_masks.append("spec.deployment_spec.env") - if deployment_spec.get("secret_env"): + if caller_provided_env_vars or deployment_spec.get("secret_env"): update_masks.append("spec.deployment_spec.secret_env") if psc_interface_config: deployment_spec["psc_interface_config"] = psc_interface_config