fix: Add pre-deployment validation for agent module imports

This change introduces a new function, _validate_agent_import, which attempts to import the user's agent.py file before the deployment process begins. This helps catch common issues such as missing dependencies, incorrect relative imports, or syntax errors in the agent code early, providing more informative error messages to the user.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 863310033

src/google/adk/cli/cli_deploy.py

@@ -14,10 +14,13 @@
 from __future__ import annotations
 
 from datetime import datetime
+import importlib
 import json
 import os
 import shutil
 import subprocess
+import sys
+import traceback
 from typing import Final
 from typing import Optional
 import warnings
@@ -465,6 +468,122 @@ 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()
+  original_sys_modules_keys = set(sys.modules.keys())
+  try:
+    # Add parent directory to path so imports work correctly
+    if parent_dir not in sys.path:
+      sys.path.insert(0, parent_dir)
+    try:
+      module = importlib.import_module(f'{module_name}.agent')
+    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 class and its dependencies are installed\n'
+            '\n'
+            'If this failure is expected (e.g., missing local dependencies), '
+            'disable agent import validation by omitting '
+            '--validate-agent-import (default) or passing '
+            '--skip-agent-import-validation (or --no-validate-agent-import).'
+        ) 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}\n\n'
+            'If this failure is expected (e.g., missing local dependencies), '
+            'disable agent import validation by omitting '
+            '--validate-agent-import (default) or passing '
+            '--skip-agent-import-validation (or --no-validate-agent-import).'
+        ) 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}\n\n'
+          'If this failure is expected (e.g., missing local dependencies), '
+          'disable agent import validation by omitting '
+          '--validate-agent-import (default) or passing '
+          '--skip-agent-import-validation (or --no-validate-agent-import).'
+      ) 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 modules introduced by validation.
+    for key in list(sys.modules.keys()):
+      if key in original_sys_modules_keys:
+        continue
+      if key == module_name or key.startswith(f'{module_name}.'):
+        sys.modules.pop(key, None)
+
+
 def _get_service_option_by_adk_version(
     adk_version: str,
     session_uri: Optional[str],
@@ -702,6 +821,7 @@ def to_agent_engine(
     requirements_file: Optional[str] = None,
     env_file: Optional[str] = None,
     agent_engine_config_file: Optional[str] = None,
+    skip_agent_import_validation: bool = True,
 ):
   """Deploys an agent to Vertex AI Agent Engine.
 
@@ -761,6 +881,11 @@ def to_agent_engine(
     agent_engine_config_file (str): The filepath to the agent engine config file
       to use. If not specified, the `.agent_engine_config.json` file in the
       `agent_folder` will be used.
+    skip_agent_import_validation (bool): Optional. Default is True. If True,
+      skip the
+      pre-deployment import validation of `agent.py`. This can be useful when
+      the local environment does not have the same dependencies as the
+      deployment environment.
   """
   app_name = os.path.basename(agent_folder)
   display_name = display_name or app_name
@@ -953,6 +1078,11 @@ 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.
+    if not skip_agent_import_validation:
+      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'

src/google/adk/cli/cli_tools_click.py

@@ -1867,6 +1867,25 @@ def cli_migrate_session(
         " directory, if any.)"
     ),
 )
+@click.option(
+    "--validate-agent-import/--no-validate-agent-import",
+    default=False,
+    help=(
+        "Optional. Validate that the agent module can be imported before"
+        " deployment. This requires your local environment to have the same"
+        " dependencies as the deployment environment. (default: disabled)"
+    ),
+)
+@click.option(
+    "--skip-agent-import-validation",
+    "skip_agent_import_validation_alias",
+    is_flag=True,
+    default=False,
+    help=(
+        "Optional. Skip pre-deployment import validation of `agent.py`. This is"
+        " the default; use --validate-agent-import to enable validation."
+    ),
+)
 @click.argument(
     "agent",
     type=click.Path(
@@ -1891,6 +1910,8 @@ def cli_deploy_agent_engine(
     requirements_file: str,
     absolutize_imports: bool,
     agent_engine_config_file: str,
+    validate_agent_import: bool = False,
+    skip_agent_import_validation_alias: bool = False,
 ):
   """Deploys an agent to Agent Engine.
 
@@ -1905,6 +1926,11 @@ def cli_deploy_agent_engine(
   """
   logging.getLogger("vertexai_genai.agentengines").setLevel(logging.INFO)
   try:
+    if validate_agent_import and skip_agent_import_validation_alias:
+      raise click.UsageError(
+          "Do not pass both --validate-agent-import and"
+          " --skip-agent-import-validation."
+      )
     cli_deploy.to_agent_engine(
         agent_folder=agent,
         project=project,
@@ -1922,6 +1948,7 @@ def cli_deploy_agent_engine(
         requirements_file=requirements_file,
         absolutize_imports=absolutize_imports,
         agent_engine_config_file=agent_engine_config_file,
+        skip_agent_import_validation=not validate_agent_import,
     )
   except Exception as e:
     click.secho(f"Deploy failed: {e}", fg="red", err=True)

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")
@@ -305,6 +306,110 @@ def test_to_agent_engine_happy_path(
   assert str(rmtree_recorder.get_last_call_args()[0]) == str(tmp_dir)
 
 
+def test_to_agent_engine_skips_agent_import_validation_by_default(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """It should skip agent.py import validation by default."""
+  validate_recorder = _Recorder()
+
+  def _validate_agent_import(*args: Any, **kwargs: Any) -> None:
+    validate_recorder(*args, **kwargs)
+    raise AssertionError("_validate_agent_import should not be called")
+
+  monkeypatch.setattr(
+      cli_deploy, "_validate_agent_import", _validate_agent_import
+  )
+
+  fake_vertexai = types.ModuleType("vertexai")
+
+  class _FakeAgentEngines:
+
+    def create(self, *, config: Dict[str, Any]) -> Any:
+      del config
+      return types.SimpleNamespace(
+          api_resource=types.SimpleNamespace(
+              name="projects/p/locations/l/reasoningEngines/e"
+          )
+      )
+
+  class _FakeVertexClient:
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+      del args
+      del kwargs
+      self.agent_engines = _FakeAgentEngines()
+
+  fake_vertexai.Client = _FakeVertexClient
+  monkeypatch.setitem(sys.modules, "vertexai", fake_vertexai)
+
+  src_dir = agent_dir(False, False)
+  cli_deploy.to_agent_engine(
+      agent_folder=str(src_dir),
+      temp_folder="tmp",
+      adk_app="my_adk_app",
+      trace_to_cloud=True,
+      project="my-gcp-project",
+      region="us-central1",
+      display_name="My Test Agent",
+      description="A test agent.",
+  )
+
+  assert validate_recorder.calls == []
+
+
+def test_to_agent_engine_validates_agent_import_when_enabled(
+    monkeypatch: pytest.MonkeyPatch,
+    agent_dir: Callable[[bool, bool], Path],
+) -> None:
+  """It should run agent.py import validation when enabled."""
+  validate_recorder = _Recorder()
+
+  def _validate_agent_import(*args: Any, **kwargs: Any) -> None:
+    validate_recorder(*args, **kwargs)
+
+  monkeypatch.setattr(
+      cli_deploy, "_validate_agent_import", _validate_agent_import
+  )
+
+  fake_vertexai = types.ModuleType("vertexai")
+
+  class _FakeAgentEngines:
+
+    def create(self, *, config: Dict[str, Any]) -> Any:
+      del config
+      return types.SimpleNamespace(
+          api_resource=types.SimpleNamespace(
+              name="projects/p/locations/l/reasoningEngines/e"
+          )
+      )
+
+  class _FakeVertexClient:
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+      del args
+      del kwargs
+      self.agent_engines = _FakeAgentEngines()
+
+  fake_vertexai.Client = _FakeVertexClient
+  monkeypatch.setitem(sys.modules, "vertexai", fake_vertexai)
+
+  src_dir = agent_dir(False, False)
+  cli_deploy.to_agent_engine(
+      agent_folder=str(src_dir),
+      temp_folder="tmp",
+      adk_app="my_adk_app",
+      trace_to_cloud=True,
+      project="my-gcp-project",
+      region="us-central1",
+      display_name="My Test Agent",
+      description="A test agent.",
+      skip_agent_import_validation=False,
+  )
+
+  assert len(validate_recorder.calls) == 1
+
+
 @pytest.mark.parametrize("include_requirements", [True, False])
 def test_to_gke_happy_path(
     monkeypatch: pytest.MonkeyPatch,
@@ -407,3 +512,145 @@ def test_to_gke_happy_path(
 
   # 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_success_with_relative_imports(self, tmp_path: Path) -> None:
+    """Should succeed when agent.py uses relative imports."""
+    (tmp_path / "helper.py").write_text("VALUE = 'my_agent'\n")
+    (tmp_path / "__init__.py").touch()
+    (tmp_path / "agent.py").write_text(
+        "from .helper import VALUE\n\nroot_agent = VALUE\n"
+    )
+
+    cli_deploy._validate_agent_import(
+        str(tmp_path), "root_agent", 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 module_name not in sys.modules
+    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 module_name not in sys.modules
+    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