diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index 47a5b688be..0dd3c6afff 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/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 diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 12edb32993..dd6281d547 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/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) diff --git a/tests/unittests/cli/utils/test_cli_deploy.py b/tests/unittests/cli/utils/test_cli_deploy.py index 98dc493687..ff5bfc618a 100644 --- a/tests/unittests/cli/utils/test_cli_deploy.py +++ b/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