feat(cli): Support custom service accounts 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,75 @@
 _AGENT_ENGINE_REQUIREMENT: Final[str] = (
     'google-cloud-aiplatform[adk,agent_engines]'
 )
+# Runtime service account email for Agent Engine, e.g.
+# my-agent@my-project.iam.gserviceaccount.com
+_SERVICE_ACCOUNT_EMAIL_RE: Final[re.Pattern[str]] = re.compile(
+    r'^[a-zA-Z0-9._+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
+)
+
+
+def _validate_service_account(service_account: str) -> str:
+  """Validates an Agent Engine runtime service account email.
+
+  Args:
+    service_account: Google Cloud service account email.
+
+  Returns:
+    The validated service account email.
+
+  Raises:
+    click.ClickException: If the email is empty or malformed.
+  """
+  service_account = service_account.strip()
+  if not service_account:
+    raise click.ClickException(
+        'service_account must be a non-empty service account email.'
+    )
+  if not _SERVICE_ACCOUNT_EMAIL_RE.fullmatch(service_account):
+    raise click.ClickException(
+        'Invalid service_account email. Expected a Google Cloud service'
+        ' account email such as'
+        ' my-agent@my-project.iam.gserviceaccount.com.'
+        f' Got: {service_account}'
+    )
+  return service_account
+
+
+def _apply_service_account_to_agent_config(
+    agent_config: dict[str, Any],
+    service_account: Optional[str],
+) -> None:
+  """Sets top-level ``service_account`` on the Agent Engine update config.
+
+  Precedence (highest last):
+
+  1. Existing ``service_account`` in ``.agent_engine_config.json``.
+  2. Explicit ``service_account`` argument (CLI / resolved from
+     ``GOOGLE_CLOUD_SERVICE_ACCOUNT``), which overrides the config file.
+
+  The Vertex Agent Engine SDK maps ``config.service_account`` onto
+  ``spec.service_account`` (runtime identity). This is distinct from
+  ``build_config.service_account`` (Cloud Build identity).
+  """
+  if service_account is not None:
+    validated = _validate_service_account(service_account)
+    existing = agent_config.get('service_account')
+    if existing and existing != validated:
+      click.echo(
+          'Overriding service_account in agent platform config with'
+          f' {validated}'
+      )
+    agent_config['service_account'] = validated
+    return
+
+  existing = agent_config.get('service_account')
+  if existing is None:
+    return
+  if not isinstance(existing, str):
+    raise click.ClickException(
+        'service_account in agent platform config must be a string email.'
+    )
+  agent_config['service_account'] = _validate_service_account(existing)
 
 
 def _ensure_agent_engine_dependency(requirements_txt_path: str) -> None:
@@ -886,6 +956,7 @@ def to_agent_engine(
     artifact_service_uri: Optional[str] = None,
     adk_version: Optional[str] = None,
     extra_packages: Optional[list[str]] = None,
+    service_account: Optional[str] = None,
 ) -> None:
   """Deploys an agent to Gemini Enterprise Agent Platform.
 
@@ -952,6 +1023,11 @@ 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.
+    service_account (str): Optional. Google Cloud service account email used
+      as the Agent Engine runtime identity. Overrides
+      ``GOOGLE_CLOUD_SERVICE_ACCOUNT`` in the ``.env`` file and
+      ``service_account`` in ``.agent_engine_config.json`` when both are
+      present. When omitted, Agent Engine uses its default service agent.
   """
   app_name = os.path.basename(agent_folder)
   display_name = display_name or app_name
@@ -1131,6 +1207,23 @@ def to_agent_engine(
           else:
             region = env_region
             click.echo(f'{region=} set by GOOGLE_CLOUD_LOCATION in {env_file}')
+      # Pop so the SA email is not forwarded as a runtime env var.
+      if 'GOOGLE_CLOUD_SERVICE_ACCOUNT' in env_vars:
+        env_service_account = env_vars.pop('GOOGLE_CLOUD_SERVICE_ACCOUNT')
+        if env_service_account:
+          if service_account:
+            click.secho(
+                'Ignoring GOOGLE_CLOUD_SERVICE_ACCOUNT in .env as'
+                ' `--service_account` was explicitly passed and takes'
+                ' precedence',
+                fg='yellow',
+            )
+          else:
+            service_account = env_service_account
+            click.echo(
+                f'{service_account=} set by GOOGLE_CLOUD_SERVICE_ACCOUNT in'
+                f' {env_file}'
+            )
     if api_key:
       if 'GOOGLE_API_KEY' in env_vars:
         click.secho(
@@ -1174,6 +1267,8 @@ def to_agent_engine(
     # Set env_vars in agent_config to None if it is not set.
     agent_config['env_vars'] = agent_config.get('env_vars', env_vars)
 
+    _apply_service_account_to_agent_config(agent_config, service_account)
+
     import vertexai
 
     from ..utils._google_client_headers import get_tracking_headers

src/google/adk/cli/cli_tools_click.py

@@ -2702,6 +2702,18 @@ def cli_migrate_session(
         " Repeatable."
     ),
 )
+@click.option(
+    "--service_account",
+    type=str,
+    default=None,
+    help=(
+        "Optional. Google Cloud service account email used as the Agent Engine"
+        " runtime identity (e.g. my-agent@my-project.iam.gserviceaccount.com)."
+        " Overrides GOOGLE_CLOUD_SERVICE_ACCOUNT in the .env file and"
+        " service_account in .agent_engine_config.json. When omitted, Agent"
+        " Engine uses its default service agent."
+    ),
+)
 @adk_services_options(default_use_local_storage=False)
 @click.argument(
     "agent",
@@ -2736,6 +2748,7 @@ def cli_deploy_agent_engine(
     session_service_uri: str | None = None,
     use_local_storage: bool = False,
     extra_packages: tuple[str, ...] = (),
+    service_account: str | None = None,
 ):
   """Deploys an agent to Agent Engine.
 
@@ -2749,6 +2762,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 custom runtime service account
+    adk deploy agent_engine --project=[project] --region=[region]
+      --service_account=my-agent@[project].iam.gserviceaccount.com
+      my_agent
   """
   logging.getLogger("vertexai_genai.agentengines").setLevel(logging.INFO)
   try:
@@ -2783,6 +2802,7 @@ def cli_deploy_agent_engine(
         session_service_uri=session_service_uri,
         adk_version=adk_version,
         extra_packages=list(extra_packages),
+        service_account=service_account,
     )
   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,193 @@ def test_to_agent_engine_extra_packages_requirements_txt_is_not_clobbered(
   assert (tmp_dir / "requirements.txt").read_text() == (
       "some-unrelated-package\n"
   )
+
+
+_VALID_SERVICE_ACCOUNT = "my-agent@my-gcp-project.iam.gserviceaccount.com"
+
+
+def test_validate_service_account_accepts_email() -> None:
+  """A well-formed service account email is accepted."""
+  assert cli_deploy._validate_service_account(_VALID_SERVICE_ACCOUNT) == (
+      _VALID_SERVICE_ACCOUNT
+  )
+
+
+@pytest.mark.parametrize(
+    "bad_sa",
+    [
+        "",
+        "   ",
+        "not-an-email",
+        "missing-domain@",
+        "@no-local-part.com",
+        "spaces emma.t@example.net",
+    ],
+)
+def test_validate_service_account_rejects_malformed_emails(
+    bad_sa: str,
+) -> None:
+  """Malformed service account emails raise a clear ClickException."""
+  with pytest.raises(click.ClickException):
+    cli_deploy._validate_service_account(bad_sa)
+
+
+def test_apply_service_account_sets_top_level_config_field() -> None:
+  """CLI service_account is set as a top-level Agent Engine config field."""
+  agent_config: Dict[str, Any] = {}
+  cli_deploy._apply_service_account_to_agent_config(
+      agent_config, _VALID_SERVICE_ACCOUNT
+  )
+  assert agent_config["service_account"] == _VALID_SERVICE_ACCOUNT
+
+
+def test_apply_service_account_cli_overrides_config_file() -> None:
+  """Explicit CLI service_account overrides the config-file value."""
+  override = "other@my-gcp-project.iam.gserviceaccount.com"
+  agent_config: Dict[str, Any] = {
+      "service_account": _VALID_SERVICE_ACCOUNT,
+  }
+  cli_deploy._apply_service_account_to_agent_config(agent_config, override)
+  assert agent_config["service_account"] == override
+
+
+def test_apply_service_account_validates_config_file_value() -> None:
+  """service_account from .agent_engine_config.json is validated in place."""
+  agent_config: Dict[str, Any] = {
+      "service_account": _VALID_SERVICE_ACCOUNT,
+  }
+  cli_deploy._apply_service_account_to_agent_config(agent_config, None)
+  assert agent_config["service_account"] == _VALID_SERVICE_ACCOUNT
+
+
+def test_to_agent_engine_forwards_service_account_in_update_config(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """to_agent_engine puts service_account on agent_engines.update config."""
+  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",
+      service_account=_VALID_SERVICE_ACCOUNT,
+  )
+
+  assert len(captured) == 1
+  assert captured[0]["service_account"] == _VALID_SERVICE_ACCOUNT
+
+
+def test_to_agent_engine_reads_service_account_from_config_file(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """service_account 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({"service_account": _VALID_SERVICE_ACCOUNT})
+  )
+
+  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]["service_account"] == _VALID_SERVICE_ACCOUNT
+
+
+def test_to_agent_engine_reads_service_account_from_env_file(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """GOOGLE_CLOUD_SERVICE_ACCOUNT in .env becomes runtime service_account."""
+  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 / ".env").write_text(
+      f"GOOGLE_CLOUD_SERVICE_ACCOUNT={_VALID_SERVICE_ACCOUNT}\n"
+      "OTHER_VAR=keep-me\n"
+  )
+
+  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]["service_account"] == _VALID_SERVICE_ACCOUNT
+  # SA email must not leak into the deployed runtime env vars.
+  assert "GOOGLE_CLOUD_SERVICE_ACCOUNT" not in captured[0]["env_vars"]
+  assert captured[0]["env_vars"]["OTHER_VAR"] == "keep-me"
+
+
+def test_to_agent_engine_rejects_invalid_service_account(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """An invalid --service_account value fails before 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",
+        service_account="not-an-email",
+    )
+
+  assert "Invalid service_account" in str(exc_info.value)
+  assert captured == []
+
+
+def test_cli_deploy_agent_engine_passes_service_account(
+    tmp_path: Path,
+) -> None:
+  """--service_account 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"--service_account={_VALID_SERVICE_ACCOUNT}",
+            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["service_account"] == _VALID_SERVICE_ACCOUNT