From 44ad42ccd50f48f76a870ec33bcb75b05cd6e515 Mon Sep 17 00:00:00 2001 From: Philippe Granger Date: Tue, 27 Jan 2026 22:13:34 +0100 Subject: [PATCH] fix(cli): add pre-deployment validation for agent module imports Add pre-deployment validation to `adk deploy agent_engine` that catches import errors before deployment. This provides clearer error messages and prevents deployments that would fail at runtime. The validation: - Checks that agent.py exists and can be imported - Verifies the expected export (root_agent or app) is present - Provides specific guidance for BaseLlm-related import errors - Properly cleans up sys.path and sys.modules after validation Note: The cloudpickle.register_pickle_by_value() workaround has been removed as Agent Engine now uses source-based deployment via source_packages, which uploads the full source code directly. This makes pickle-based serialization unnecessary for custom classes. Fixes #4208 --- src/google/adk/cli/cli_deploy.py | 124 +++++++++++++++++ tests/unittests/cli/utils/test_cli_deploy.py | 133 ++++++++++++++++++- 2 files changed, 255 insertions(+), 2 deletions(-) diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index d6394ff2fd..47175503fc 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -14,10 +14,14 @@ from __future__ import annotations from datetime import datetime +import importlib.util import json +import logging import os import shutil import subprocess +import sys +import traceback from typing import Final from typing import Optional import warnings @@ -25,6 +29,8 @@ import click from packaging.version import parse +logger = logging.getLogger('google_adk.' + __name__) + _IS_WINDOWS = os.name == 'nt' _GCLOUD_CMD = 'gcloud.cmd' if _IS_WINDOWS else 'gcloud' _LOCAL_STORAGE_FLAG_MIN_VERSION: Final[str] = '1.21.0' @@ -465,6 +471,120 @@ def _validate_gcloud_extra_args( ) +def _validate_agent_import( + agent_src_path: str, + adk_app_object: str, + is_config_agent: bool, +) -> None: + """Validates that the agent module can be imported successfully. + + This pre-deployment validation catches common issues like missing + dependencies or import errors in custom BaseLlm implementations before + the agent is deployed to Agent Engine. This provides clearer error + messages and prevents deployments that would fail at runtime. + + Args: + agent_src_path: Path to the staged agent source code. + adk_app_object: The Python object name to import ('root_agent' or 'app'). + is_config_agent: Whether this is a config-based agent. + + Raises: + click.ClickException: If the agent module cannot be imported. + """ + if is_config_agent: + # Config agents are loaded from YAML, skip Python import validation + return + + agent_module_path = os.path.join(agent_src_path, 'agent.py') + if not os.path.exists(agent_module_path): + raise click.ClickException( + f'Agent module not found at {agent_module_path}. ' + 'Please ensure your agent folder contains an agent.py file.' + ) + + # Add the parent directory to sys.path temporarily for import resolution + parent_dir = os.path.dirname(agent_src_path) + module_name = os.path.basename(agent_src_path) + + original_sys_path = sys.path.copy() + try: + # Add parent directory to path so imports work correctly + if parent_dir not in sys.path: + sys.path.insert(0, parent_dir) + + # Load the agent module spec + spec = importlib.util.spec_from_file_location( + f'{module_name}.agent', + agent_module_path, + submodule_search_locations=[agent_src_path], + ) + if spec is None or spec.loader is None: + raise click.ClickException( + f'Failed to load module spec from {agent_module_path}' + ) + + # Try to load the module + module = importlib.util.module_from_spec(spec) + sys.modules[f'{module_name}.agent'] = module + + try: + spec.loader.exec_module(module) + except ImportError as e: + error_msg = str(e) + tb = traceback.format_exc() + + # Check for common issues + if 'BaseLlm' in tb or 'base_llm' in tb.lower(): + raise click.ClickException( + 'Failed to import agent module due to a BaseLlm-related error:\n' + f'{error_msg}\n\n' + 'This error often occurs when deploying agents with custom LLM ' + 'implementations. Please ensure:\n' + '1. All custom LLM classes are defined in files within your agent ' + 'folder\n' + '2. All required dependencies are listed in requirements.txt\n' + '3. Import paths use relative imports (e.g., "from .my_llm import ' + 'MyLlm")\n' + '4. Your custom BaseLlm implementation is serializable' + ) from e + else: + raise click.ClickException( + f'Failed to import agent module:\n{error_msg}\n\n' + 'Please ensure all dependencies are listed in requirements.txt ' + 'and all imports are resolvable.\n\n' + f'Full traceback:\n{tb}' + ) from e + except Exception as e: + tb = traceback.format_exc() + raise click.ClickException( + f'Error while loading agent module:\n{e}\n\n' + 'Please check your agent code for errors.\n\n' + f'Full traceback:\n{tb}' + ) from e + + # Check that the expected object exists + if not hasattr(module, adk_app_object): + available_attrs = [ + attr for attr in dir(module) if not attr.startswith('_') + ] + raise click.ClickException( + f"Agent module does not export '{adk_app_object}'. " + f'Available exports: {available_attrs}\n\n' + 'Please ensure your agent.py exports either "root_agent" or "app".' + ) + + click.echo( + 'Agent module validation successful: ' + f'found "{adk_app_object}" in agent.py' + ) + + finally: + # Restore original sys.path + sys.path[:] = original_sys_path + # Clean up the module from sys.modules + sys.modules.pop(f'{module_name}.agent', None) + + def _get_service_option_by_adk_version( adk_version: str, session_uri: Optional[str], @@ -953,6 +1073,10 @@ def to_agent_engine( click.echo(f'Config agent detected: {config_root_agent_file}') is_config_agent = True + # Validate that the agent module can be imported before deployment + click.echo('Validating agent module...') + _validate_agent_import(agent_src_path, adk_app_object, is_config_agent) + adk_app_file = os.path.join(temp_folder, f'{adk_app}.py') if adk_app_object == 'root_agent': adk_app_type = 'agent' diff --git a/tests/unittests/cli/utils/test_cli_deploy.py b/tests/unittests/cli/utils/test_cli_deploy.py index 36484da149..805b7ebc2f 100644 --- a/tests/unittests/cli/utils/test_cli_deploy.py +++ b/tests/unittests/cli/utils/test_cli_deploy.py @@ -14,7 +14,6 @@ """Tests for utilities in cli_deploy.""" - from __future__ import annotations import importlib @@ -83,7 +82,9 @@ def agent_dir(tmp_path: Path) -> Callable[[bool, bool], Path]: def _factory(include_requirements: bool, include_env: bool) -> Path: base = tmp_path / "agent" base.mkdir() - (base / "agent.py").write_text("# dummy agent") + (base / "agent.py").write_text( + "# dummy agent\nroot_agent = 'dummy_agent'\n" + ) (base / "__init__.py").touch() if include_requirements: (base / "requirements.txt").write_text("pytest\n") @@ -407,3 +408,131 @@ def mock_subprocess_run(*args, **kwargs): # 4. Verify cleanup assert str(rmtree_recorder.get_last_call_args()[0]) == str(tmp_path) + + +# _validate_agent_import tests +class TestValidateAgentImport: + """Tests for the _validate_agent_import function.""" + + def test_skips_config_agents(self, tmp_path: Path) -> None: + """Config agents should skip validation.""" + # This should not raise even with no agent.py file + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=True + ) + + def test_raises_on_missing_agent_module(self, tmp_path: Path) -> None: + """Should raise when agent.py is missing.""" + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "Agent module not found" in str(exc_info.value) + + def test_raises_on_missing_export(self, tmp_path: Path) -> None: + """Should raise when the expected export is missing.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("some_other_var = 'hello'\n") + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "does not export 'root_agent'" in str(exc_info.value) + assert "some_other_var" in str(exc_info.value) + + def test_success_with_root_agent_export(self, tmp_path: Path) -> None: + """Should succeed when root_agent is exported.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("root_agent = 'my_agent'\n") + (tmp_path / "__init__.py").touch() + + # Should not raise + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + + def test_success_with_app_export(self, tmp_path: Path) -> None: + """Should succeed when app is exported.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("app = 'my_app'\n") + (tmp_path / "__init__.py").touch() + + # Should not raise + cli_deploy._validate_agent_import( + str(tmp_path), "app", is_config_agent=False + ) + + def test_raises_on_import_error(self, tmp_path: Path) -> None: + """Should raise with helpful message on ImportError.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("from nonexistent_module import something\n") + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "Failed to import agent module" in str(exc_info.value) + assert "nonexistent_module" in str(exc_info.value) + + def test_raises_on_basellm_import_error(self, tmp_path: Path) -> None: + """Should provide specific guidance for BaseLlm import errors.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text( + "from google.adk.models.base_llm import NonexistentBaseLlm\n" + ) + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "BaseLlm-related error" in str(exc_info.value) + assert "custom LLM" in str(exc_info.value) + + def test_raises_on_syntax_error(self, tmp_path: Path) -> None: + """Should raise on syntax errors in agent.py.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("def invalid syntax here:\n") + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "Error while loading agent module" in str(exc_info.value) + + def test_cleans_up_sys_modules(self, tmp_path: Path) -> None: + """Should clean up sys.modules after validation.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("root_agent = 'my_agent'\n") + (tmp_path / "__init__.py").touch() + + module_name = tmp_path.name + agent_module_key = f"{module_name}.agent" + + # Ensure module is not in sys.modules before + assert agent_module_key not in sys.modules + + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + + # Ensure module is cleaned up after + assert agent_module_key not in sys.modules + + def test_restores_sys_path(self, tmp_path: Path) -> None: + """Should restore sys.path after validation.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("root_agent = 'my_agent'\n") + (tmp_path / "__init__.py").touch() + + original_path = sys.path.copy() + + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + + assert sys.path == original_path