feat(cli): Support Cloud Build worker pools for Agent Engine deploy

src/google/adk/cli/cli_deploy.py

@@ -17,6 +17,7 @@
 import importlib
 import json
 import os
+import re
 import shutil
 import subprocess
 import sys
@@ -40,6 +41,92 @@
 _AGENT_ENGINE_REQUIREMENT: Final[str] = (
     'google-cloud-aiplatform[adk,agent_engines]'
 )
+# Full Cloud Build private worker pool resource name, e.g.
+# projects/my-project/locations/us-central1/workerPools/my-private-pool
+_WORKER_POOL_RESOURCE_RE: Final[re.Pattern[str]] = re.compile(
+    r'^projects/[^/]+/locations/[^/]+/workerPools/[^/]+$'
+)
+
+
+def _validate_worker_pool(worker_pool: str) -> str:
+  """Validates a Cloud Build worker pool resource name.
+
+  Args:
+    worker_pool: Full resource name of the form
+      `projects/{project}/locations/{location}/workerPools/{pool}`.
+
+  Returns:
+    The validated worker pool resource name.
+
+  Raises:
+    click.ClickException: If the resource name is empty or malformed.
+  """
+  worker_pool = worker_pool.strip()
+  if not worker_pool:
+    raise click.ClickException('worker_pool must be a non-empty resource name.')
+  if not _WORKER_POOL_RESOURCE_RE.fullmatch(worker_pool):
+    raise click.ClickException(
+        'Invalid worker_pool resource name. Expected format:'
+        ' projects/{project}/locations/{location}/workerPools/{pool}.'
+        f' Got: {worker_pool}'
+    )
+  return worker_pool
+
+
+def _apply_worker_pool_to_agent_config(
+    agent_config: dict[str, Any],
+    worker_pool: Optional[str],
+) -> None:
+  """Nests worker_pool into agent_config['build_config'].
+
+  Supports three sources, in increasing precedence:
+
+  1. Existing ``build_config.worker_pool`` already in ``agent_config``.
+  2. Top-level ``worker_pool`` convenience key in ``.agent_engine_config.json``
+     (popped so it is not forwarded as an unknown top-level field).
+  3. Explicit ``worker_pool`` argument (CLI flag), which overrides both.
+
+  The Vertex Agent Engine SDK reads Cloud Build private pools from
+  ``config.build_config.worker_pool`` and maps them onto
+  ``spec.build_spec.worker_pool``.
+  """
+  build_config = agent_config.get('build_config')
+  if build_config is None:
+    build_config = {}
+  elif not isinstance(build_config, dict):
+    raise click.ClickException(
+        'build_config in agent platform config must be a JSON object.'
+    )
+  else:
+    # Copy so we do not mutate a shared structure unexpectedly.
+    build_config = dict(build_config)
+
+  config_worker_pool = agent_config.pop('worker_pool', None)
+  if config_worker_pool is not None:
+    if not isinstance(config_worker_pool, str):
+      raise click.ClickException(
+          'worker_pool in agent platform config must be a string resource name.'
+      )
+    build_config['worker_pool'] = _validate_worker_pool(config_worker_pool)
+
+  if worker_pool is not None:
+    if build_config.get('worker_pool'):
+      click.echo(
+          'Overriding build_config.worker_pool in agent platform config with'
+          f' {worker_pool}'
+      )
+    build_config['worker_pool'] = _validate_worker_pool(worker_pool)
+
+  # Validate any worker_pool that was already nested under build_config.
+  if 'worker_pool' in build_config and build_config['worker_pool'] is not None:
+    build_config['worker_pool'] = _validate_worker_pool(
+        str(build_config['worker_pool'])
+    )
+
+  if build_config:
+    agent_config['build_config'] = build_config
+  else:
+    agent_config.pop('build_config', None)
 
 
 def _ensure_agent_engine_dependency(requirements_txt_path: str) -> None:
@@ -886,6 +973,7 @@ def to_agent_engine(
     artifact_service_uri: Optional[str] = None,
     adk_version: Optional[str] = None,
     extra_packages: Optional[list[str]] = None,
+    worker_pool: Optional[str] = None,
 ) -> None:
   """Deploys an agent to Gemini Enterprise Agent Platform.
 
@@ -952,6 +1040,13 @@ def to_agent_engine(
       used.
     extra_packages (list[str]): Optional. Additional local file or directory
       paths to stage alongside the agent and make importable in the image.
+    worker_pool (str): Optional. Full Cloud Build private worker pool resource
+      name
+      (`projects/{project}/locations/{location}/workerPools/{pool}`).
+      When set, Agent Engine builds the container image on that pool so
+      deploys can reach private networks / comply with org build policies.
+      Overrides `worker_pool` / `build_config.worker_pool` from
+      `.agent_engine_config.json` when both are present.
   """
   app_name = os.path.basename(agent_folder)
   display_name = display_name or app_name
@@ -1049,6 +1144,8 @@ def to_agent_engine(
         )
       agent_config['description'] = description
 
+    _apply_worker_pool_to_agent_config(agent_config, worker_pool)
+
     config_extra_packages = agent_config.pop('extra_packages', None) or []
     # CLI entries resolve against the invocation dir; config-file entries
     # against the agent folder that declared them.

src/google/adk/cli/cli_tools_click.py

@@ -2702,6 +2702,19 @@ def cli_migrate_session(
         " Repeatable."
     ),
 )
+@click.option(
+    "--worker_pool",
+    type=str,
+    default=None,
+    help=(
+        "Optional. Cloud Build private worker pool resource name used to build"
+        " the Agent Engine container image. Format:"
+        " projects/{project}/locations/{location}/workerPools/{pool}."
+        " Required for VPC-SC / private-network environments that cannot use"
+        " the default public Cloud Build pool. Overrides `worker_pool` or"
+        " `build_config.worker_pool` in `.agent_engine_config.json`."
+    ),
+)
 @adk_services_options(default_use_local_storage=False)
 @click.argument(
     "agent",
@@ -2736,6 +2749,7 @@ def cli_deploy_agent_engine(
     session_service_uri: str | None = None,
     use_local_storage: bool = False,
     extra_packages: tuple[str, ...] = (),
+    worker_pool: str | None = None,
 ):
   """Deploys an agent to Agent Engine.
 
@@ -2749,6 +2763,12 @@ def cli_deploy_agent_engine(
     # With Google Cloud Project and Region
     adk deploy agent_engine --project=[project] --region=[region]
       --display_name=[app_name] my_agent
+
+    \b
+    # With a private Cloud Build worker pool (VPC-SC / private network)
+    adk deploy agent_engine --project=[project] --region=[region]
+      --worker_pool=projects/[project]/locations/[region]/workerPools/[pool]
+      my_agent
   """
   logging.getLogger("vertexai_genai.agentengines").setLevel(logging.INFO)
   try:
@@ -2783,6 +2803,7 @@ def cli_deploy_agent_engine(
         session_service_uri=session_service_uri,
         adk_version=adk_version,
         extra_packages=list(extra_packages),
+        worker_pool=worker_pool,
     )
   except Exception as e:
     click.secho(f"Deploy failed: {e}", fg="red", err=True)

tests/unittests/cli/utils/test_cli_deploy.py

@@ -1164,3 +1164,178 @@ def test_to_agent_engine_extra_packages_requirements_txt_is_not_clobbered(
   assert (tmp_dir / "requirements.txt").read_text() == (
       "some-unrelated-package\n"
   )
+
+
+_VALID_WORKER_POOL = (
+    "projects/my-gcp-project/locations/us-central1/workerPools/my-private-pool"
+)
+
+
+def test_validate_worker_pool_accepts_full_resource_name() -> None:
+  """A well-formed Cloud Build worker pool resource name is accepted."""
+  assert cli_deploy._validate_worker_pool(_VALID_WORKER_POOL) == (
+      _VALID_WORKER_POOL
+  )
+
+
+@pytest.mark.parametrize(
+    "bad_pool",
+    [
+        "",
+        "   ",
+        "my-private-pool",
+        "projects/p/locations/l/workerPools/",
+        "projects/p/locations/l/pools/my-pool",
+        "projects/p/workerPools/my-pool",
+    ],
+)
+def test_validate_worker_pool_rejects_malformed_names(bad_pool: str) -> None:
+  """Malformed worker pool resource names raise a clear ClickException."""
+  with pytest.raises(click.ClickException):
+    cli_deploy._validate_worker_pool(bad_pool)
+
+
+def test_apply_worker_pool_nests_cli_value_into_build_config() -> None:
+  """CLI worker_pool is nested under build_config for the Vertex SDK."""
+  agent_config: Dict[str, Any] = {}
+  cli_deploy._apply_worker_pool_to_agent_config(
+      agent_config, _VALID_WORKER_POOL
+  )
+  assert agent_config["build_config"]["worker_pool"] == _VALID_WORKER_POOL
+  assert "worker_pool" not in agent_config
+
+
+def test_apply_worker_pool_pops_top_level_config_key() -> None:
+  """Top-level worker_pool in .agent_engine_config.json is nested and removed."""
+  agent_config: Dict[str, Any] = {"worker_pool": _VALID_WORKER_POOL}
+  cli_deploy._apply_worker_pool_to_agent_config(agent_config, None)
+  assert agent_config["build_config"]["worker_pool"] == _VALID_WORKER_POOL
+  assert "worker_pool" not in agent_config
+
+
+def test_apply_worker_pool_cli_overrides_config_file() -> None:
+  """Explicit CLI worker_pool overrides values from the config file."""
+  override = "projects/other/locations/europe-west1/workerPools/compliance-pool"
+  agent_config: Dict[str, Any] = {
+      "build_config": {"worker_pool": _VALID_WORKER_POOL},
+  }
+  cli_deploy._apply_worker_pool_to_agent_config(agent_config, override)
+  assert agent_config["build_config"]["worker_pool"] == override
+
+
+def test_apply_worker_pool_preserves_other_build_config_fields() -> None:
+  """Existing build_config.service_account is kept when adding worker_pool."""
+  agent_config: Dict[str, Any] = {
+      "build_config": {
+          "service_account": "builder@example.iam.gserviceaccount.com",
+      },
+  }
+  cli_deploy._apply_worker_pool_to_agent_config(
+      agent_config, _VALID_WORKER_POOL
+  )
+  assert agent_config["build_config"] == {
+      "service_account": "builder@example.iam.gserviceaccount.com",
+      "worker_pool": _VALID_WORKER_POOL,
+  }
+
+
+def test_to_agent_engine_forwards_worker_pool_in_update_config(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """to_agent_engine puts worker_pool under build_config on agent_engines.update."""
+  monkeypatch.setattr(shutil, "rmtree", _Recorder())
+  captured: List[Dict[str, Any]] = []
+  monkeypatch.setitem(
+      sys.modules, "vertexai", _make_recording_vertexai(captured)
+  )
+  src_dir = agent_dir(False, False)
+
+  cli_deploy.to_agent_engine(
+      agent_folder=str(src_dir),
+      temp_folder="tmp",
+      project="my-gcp-project",
+      region="us-central1",
+      adk_version="1.2.0",
+      worker_pool=_VALID_WORKER_POOL,
+  )
+
+  assert len(captured) == 1
+  assert captured[0]["build_config"]["worker_pool"] == _VALID_WORKER_POOL
+
+
+def test_to_agent_engine_reads_worker_pool_from_config_file(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """worker_pool from .agent_engine_config.json is forwarded on deploy."""
+  monkeypatch.setattr(shutil, "rmtree", _Recorder())
+  captured: List[Dict[str, Any]] = []
+  monkeypatch.setitem(
+      sys.modules, "vertexai", _make_recording_vertexai(captured)
+  )
+  src_dir = agent_dir(False, False)
+  (src_dir / ".agent_engine_config.json").write_text(
+      json.dumps({"worker_pool": _VALID_WORKER_POOL})
+  )
+
+  cli_deploy.to_agent_engine(
+      agent_folder=str(src_dir),
+      temp_folder="tmp",
+      project="my-gcp-project",
+      region="us-central1",
+      adk_version="1.2.0",
+  )
+
+  assert captured[0]["build_config"]["worker_pool"] == _VALID_WORKER_POOL
+  assert "worker_pool" not in captured[0]
+
+
+def test_to_agent_engine_rejects_invalid_worker_pool(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """An invalid --worker_pool value fails before calling Agent Engine APIs."""
+  monkeypatch.setattr(shutil, "rmtree", _Recorder())
+  captured: List[Dict[str, Any]] = []
+  monkeypatch.setitem(
+      sys.modules, "vertexai", _make_recording_vertexai(captured)
+  )
+  src_dir = agent_dir(False, False)
+
+  with pytest.raises(click.ClickException) as exc_info:
+    cli_deploy.to_agent_engine(
+        agent_folder=str(src_dir),
+        temp_folder="tmp",
+        project="my-gcp-project",
+        region="us-central1",
+        adk_version="1.2.0",
+        worker_pool="not-a-resource-name",
+    )
+
+  assert "Invalid worker_pool" in str(exc_info.value)
+  assert captured == []
+
+
+def test_cli_deploy_agent_engine_passes_worker_pool(tmp_path: Path) -> None:
+  """--worker_pool reaches to_agent_engine as a keyword argument."""
+  agent_dir = tmp_path / "my_agent"
+  agent_dir.mkdir()
+  runner = CliRunner()
+  with mock.patch(
+      "src.google.adk.cli.cli_deploy.to_agent_engine"
+  ) as mock_to_agent_engine:
+    result = runner.invoke(
+        cli_tools_click.main,
+        [
+            "deploy",
+            "agent_engine",
+            f"--worker_pool={_VALID_WORKER_POOL}",
+            str(agent_dir),
+        ],
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    mock_to_agent_engine.assert_called_once()
+    _, kwargs = mock_to_agent_engine.call_args
+    assert kwargs["worker_pool"] == _VALID_WORKER_POOL