Merge branch 'main' into feat/lifecycle-error-callbacks-4774

Author: rohityan

pyproject.toml

@@ -111,7 +111,7 @@ eval = [
   # go/keep-sorted start
   "Jinja2>=3.1.4,<4.0.0",                         # For eval template rendering
   "gepa>=0.1.0",
-  "google-cloud-aiplatform[evaluation]>=1.140.0",
+  "google-cloud-aiplatform[evaluation]>=1.143.0",
   "pandas>=2.2.3",
   "rouge-score>=0.1.2",
   "tabulate>=0.9.0",

src/google/adk/cli/cli_deploy.py

@@ -22,6 +22,7 @@
 import sys
 import traceback
 from typing import Final
+from typing import Literal
 from typing import Optional
 import warnings
 
@@ -1170,6 +1171,9 @@ def to_gke(
     memory_service_uri: Optional[str] = None,
     use_local_storage: bool = False,
     a2a: bool = False,
+    service_type: Literal[
+        'ClusterIP', 'NodePort', 'LoadBalancer'
+    ] = 'ClusterIP',
 ):
   """Deploys an agent to Google Kubernetes Engine(GKE).
 
@@ -1197,6 +1201,7 @@ def to_gke(
     artifact_service_uri: The URI of the artifact service.
     memory_service_uri: The URI of the memory service.
     use_local_storage: Whether to use local .adk storage in the container.
+    service_type: The Kubernetes Service type (default: ClusterIP).
   """
   click.secho(
       '\n🚀 Starting ADK Agent Deployment to GKE...', fg='cyan', bold=True
@@ -1334,7 +1339,7 @@ def to_gke(
 metadata:
   name: {service_name}
 spec:
-  type: LoadBalancer
+  type: {service_type}
   selector:
     app: {service_name}
   ports:
@@ -1388,3 +1393,11 @@ def to_gke(
   click.secho(
       '\n🎉 Deployment to GKE finished successfully!', fg='cyan', bold=True
   )
+  if service_type == 'ClusterIP':
+    click.echo(
+        '\nThe service is only reachable from within the cluster.'
+        ' To access it locally, run:'
+        f'\n  kubectl port-forward svc/{service_name} {port}:{port}'
+        '\n\nTo expose the service externally, add a Gateway or'
+        ' re-deploy with --service_type=LoadBalancer.'
+    )

src/google/adk/cli/cli_tools_click.py

@@ -2238,6 +2238,17 @@ def cli_deploy_agent_engine(
     default="INFO",
     help="Optional. Set the logging level",
 )
+@click.option(
+    "--service_type",
+    type=click.Choice(["ClusterIP", "LoadBalancer"], case_sensitive=True),
+    default="ClusterIP",
+    show_default=True,
+    help=(
+        "Optional. The Kubernetes Service type for the deployed agent."
+        " ClusterIP (default) keeps the service cluster-internal;"
+        " use LoadBalancer to expose a public IP."
+    ),
+)
 @click.option(
     "--temp_folder",
     type=str,
@@ -2281,6 +2292,7 @@ def cli_deploy_gke(
     otel_to_cloud: bool,
     with_ui: bool,
     adk_version: str,
+    service_type: str,
     log_level: Optional[str] = None,
     session_service_uri: Optional[str] = None,
     artifact_service_uri: Optional[str] = None,
@@ -2312,6 +2324,7 @@ def cli_deploy_gke(
         with_ui=with_ui,
         log_level=log_level,
         adk_version=adk_version,
+        service_type=service_type,
         session_service_uri=session_service_uri,
         artifact_service_uri=artifact_service_uri,
         memory_service_uri=memory_service_uri,

src/google/adk/cli/fast_api.py

@@ -291,7 +291,7 @@ def _normalize_relative_path(path: str) -> str:
     def _has_parent_reference(path: str) -> bool:
       return any(part == ".." for part in path.split("/"))
 
-    _ALLOWED_UPLOAD_EXTENSIONS = frozenset({".yaml", ".yml"})
+    _ALLOWED_EXTENSIONS = frozenset({".yaml", ".yml"})
 
     def _parse_upload_filename(filename: Optional[str]) -> tuple[str, str]:
       if not filename:
@@ -307,10 +307,10 @@ def _parse_upload_filename(filename: Optional[str]) -> tuple[str, str]:
       if _has_parent_reference(rel_path):
         raise ValueError(f"Path traversal rejected: {filename!r}")
       ext = os.path.splitext(rel_path)[1].lower()
-      if ext not in _ALLOWED_UPLOAD_EXTENSIONS:
+      if ext not in _ALLOWED_EXTENSIONS:
         raise ValueError(
             f"File type not allowed: {rel_path!r}"
-            f" (allowed: {', '.join(sorted(_ALLOWED_UPLOAD_EXTENSIONS))})"
+            f" (allowed: {', '.join(sorted(_ALLOWED_EXTENSIONS))})"
         )
       return app_name, rel_path
 
@@ -322,6 +322,12 @@ def _parse_file_path(file_path: str) -> str:
         raise ValueError(f"Absolute file_path rejected: {file_path!r}")
       if _has_parent_reference(file_path):
         raise ValueError(f"Path traversal rejected: {file_path!r}")
+      ext = os.path.splitext(file_path)[1].lower()
+      if ext not in _ALLOWED_EXTENSIONS:
+        raise ValueError(
+            f"File type not allowed: {file_path!r}"
+            f" (allowed: {', '.join(sorted(_ALLOWED_EXTENSIONS))})"
+        )
       return file_path
 
     def _resolve_under_dir(root_dir: Path, rel_path: str) -> Path:

src/google/adk/cli/utils/agent_loader.py

@@ -19,6 +19,7 @@
 import logging
 import os
 from pathlib import Path
+import re
 import sys
 from typing import Any
 from typing import Literal
@@ -187,8 +188,36 @@ def _load_from_yaml_config(
       ) + e.args[1:]
       raise e
 
+  _VALID_AGENT_NAME_RE = re.compile(r"^[a-zA-Z0-9_]+$")
+
+  def _validate_agent_name(self, agent_name: str) -> None:
+    """Validate agent name to prevent arbitrary module imports."""
+    # Strip the special agent prefix for validation
+    if agent_name.startswith("__"):
+      name_to_check = agent_name[2:]
+      check_dir = os.path.abspath(SPECIAL_AGENTS_DIR)
+    else:
+      name_to_check = agent_name
+      check_dir = self.agents_dir
+
+    if not self._VALID_AGENT_NAME_RE.match(name_to_check):
+      raise ValueError(
+          f"Invalid agent name: {agent_name!r}. Agent names must be valid"
+          " Python identifiers (letters, digits, and underscores only)."
+      )
+
+    # Verify the agent exists on disk before allowing import
+    agent_path = Path(check_dir) / name_to_check
+    agent_file = Path(check_dir) / f"{name_to_check}.py"
+    if not (agent_path.is_dir() or agent_file.is_file()):
+      raise ValueError(
+          f"Agent not found: {agent_name!r}. No matching directory or module"
+          f" exists in '{os.path.join(check_dir, name_to_check)}'."
+      )
+
   def _perform_load(self, agent_name: str) -> Union[BaseAgent, App]:
     """Internal logic to load an agent"""
+    self._validate_agent_name(agent_name)
     # Determine the directory to use for loading
     if agent_name.startswith("__"):
       # Special agent: use special agents directory
@@ -335,18 +364,13 @@ def load_agent(self, agent_name: str) -> Union[BaseAgent, App]:
   def list_agents(self) -> list[str]:
     """Lists all agents available in the agent loader (sorted alphabetically)."""
     base_path = Path.cwd() / self.agents_dir
-    agent_names = []
-    for x in os.listdir(base_path):
-      if (
-          os.path.isdir(os.path.join(base_path, x))
-          and not x.startswith(".")
-          and x != "__pycache__"
-      ):
-        try:
-          self._determine_agent_language(x)
-          agent_names.append(x)
-        except ValueError:
-          continue
+    agent_names = [
+        x
+        for x in os.listdir(base_path)
+        if os.path.isdir(os.path.join(base_path, x))
+        and not x.startswith(".")
+        and x != "__pycache__"
+    ]
     agent_names.sort()
     return agent_names
 

src/google/adk/events/event.py

@@ -36,7 +36,7 @@ class Event(LlmResponse):
   """
 
   model_config = ConfigDict(
-      extra='forbid',
+      extra='ignore',
       ser_json_bytes='base64',
       val_json_bytes='base64',
       alias_generator=alias_generators.to_camel,

src/google/adk/sessions/database_session_service.py

@@ -33,6 +33,7 @@
 from sqlalchemy.engine import Connection
 from sqlalchemy.engine import make_url
 from sqlalchemy.exc import ArgumentError
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import async_sessionmaker
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory
@@ -103,6 +104,35 @@ async def _select_required_state(
   return state_row
 
 
+async def _get_or_create_state(
+    *,
+    sql_session: DatabaseSessionFactory,
+    state_model: type[_StorageStateT],
+    primary_key: Any,
+    defaults: dict[str, Any],
+) -> _StorageStateT:
+  """Returns an existing state row or creates one, handling concurrent inserts.
+
+  Uses a SAVEPOINT so that an IntegrityError from a racing INSERT does not
+  invalidate the outer transaction.
+  """
+  row = await sql_session.get(state_model, primary_key)
+  if row is not None:
+    return row
+  try:
+    async with sql_session.begin_nested():
+      row = state_model(**defaults)
+      sql_session.add(row)
+    return row
+  except IntegrityError:
+    # Another concurrent caller inserted the row first.
+    # The savepoint was rolled back, so re-fetch the winner's row.
+    row = await sql_session.get(state_model, primary_key)
+    if row is None:
+      raise
+    return row
+
+
 def _set_sqlite_pragma(dbapi_connection, connection_record):
   cursor = dbapi_connection.cursor()
   cursor.execute("PRAGMA foreign_keys=ON")
@@ -401,24 +431,20 @@ async def create_session(
         raise AlreadyExistsError(
             f"Session with id {session_id} already exists."
         )
-      # Fetch app and user states from storage
-      storage_app_state = await sql_session.get(
-          schema.StorageAppState, (app_name)
+      # Get or create state rows, handling concurrent insert races.
+      storage_app_state = await _get_or_create_state(
+          sql_session=sql_session,
+          state_model=schema.StorageAppState,
+          primary_key=app_name,
+          defaults={"app_name": app_name, "state": {}},
       )
-      storage_user_state = await sql_session.get(
-          schema.StorageUserState, (app_name, user_id)
+      storage_user_state = await _get_or_create_state(
+          sql_session=sql_session,
+          state_model=schema.StorageUserState,
+          primary_key=(app_name, user_id),
+          defaults={"app_name": app_name, "user_id": user_id, "state": {}},
       )
 
-      # Create state tables if not exist
-      if not storage_app_state:
-        storage_app_state = schema.StorageAppState(app_name=app_name, state={})
-        sql_session.add(storage_app_state)
-      if not storage_user_state:
-        storage_user_state = schema.StorageUserState(
-            app_name=app_name, user_id=user_id, state={}
-        )
-        sql_session.add(storage_user_state)
-
       # Extract state deltas
       state_deltas = _session_util.extract_state_delta(state)
       app_state_delta = state_deltas["app"]

src/google/adk/skills/_utils.py

@@ -20,6 +20,7 @@
 import pathlib
 from typing import Union
 
+from google.auth import credentials as auth
 from google.cloud import storage
 from pydantic import ValidationError
 import yaml
@@ -301,6 +302,8 @@ def _list_skills_in_dir(
 def _list_skills_in_gcs_dir(
     bucket_name: str,
     skills_base_path: str = "",
+    project_id: str | None = None,
+    credentials: auth.Credentials | None = None,
 ) -> Dict[str, models.Frontmatter]:
   """List skills in a GCS directory.
 
@@ -311,7 +314,7 @@ def _list_skills_in_gcs_dir(
   Returns:
     Dictionary mapping skill IDs to their frontmatter.
   """
-  client = storage.Client()
+  client = storage.Client(project=project_id, credentials=credentials)
   bucket = client.bucket(bucket_name)
 
   base_prefix = skills_base_path.strip("/")
@@ -350,13 +353,17 @@ def _load_skill_from_gcs_dir(
     bucket_name: str,
     skill_id: str,
     skills_base_path: str = "",
+    project_id: str | None = None,
+    credentials: auth.Credentials | None = None,
 ) -> models.Skill:
   """Load a complete skill from a GCS directory.
 
   Args:
     bucket_name: Name of the GCS bucket.
     skill_id: The ID of the skill (directory name).
     skills_base_path: Base directory within the bucket (e.g., 'path/to/skills').
+    project_id: Project ID to use for GCS client.
+    credentials: Credentials to use for GCS client.
 
   Returns:
     Skill object with all components loaded.
@@ -366,7 +373,8 @@ def _load_skill_from_gcs_dir(
     ValueError: If SKILL.md is invalid or the skill name does not match
       the directory name.
   """
-  client = storage.Client()
+
+  client = storage.Client(project=project_id, credentials=credentials)
   bucket = client.bucket(bucket_name)
 
   base_prefix = skills_base_path.strip("/")

src/google/adk/tools/bash_tool.py

@@ -141,8 +141,12 @@ async def run_async(
     try:
       stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=30)
       return {
-          "stdout": stdout.decode(),
-          "stderr": stderr.decode(),
+          "stdout": (
+              stdout.decode() if stdout is not None else "<No stdout captured>"
+          ),
+          "stderr": (
+              stderr.decode() if stderr is not None else "<No stderr captured>"
+          ),
           "returncode": process.returncode,
       }
     except asyncio.TimeoutError:

tests/unittests/cli/test_fast_api.py

@@ -1759,6 +1759,40 @@ def test_builder_save_allows_yaml_files(builder_test_client, tmp_path):
   assert response.json() is True
 
 
+def test_builder_get_rejects_non_yaml_file_paths(builder_test_client, tmp_path):
+  """GET /builder/app/{app_name}?file_path=... rejects non-YAML extensions."""
+  app_root = tmp_path / "app"
+  app_root.mkdir(parents=True, exist_ok=True)
+  (app_root / ".env").write_text("SECRET=supersecret\n")
+  (app_root / "agent.py").write_text("root_agent = None\n")
+  (app_root / "config.json").write_text("{}\n")
+
+  for file_path in [".env", "agent.py", "config.json"]:
+    response = builder_test_client.get(
+        f"/builder/app/app?file_path={file_path}"
+    )
+    assert response.status_code == 200, f"Expected 200 for {file_path}"
+    assert response.text == "", f"Expected empty response for {file_path}"
+
+
+def test_builder_get_allows_yaml_file_paths(builder_test_client, tmp_path):
+  """GET /builder/app/{app_name}?file_path=... allows YAML extensions."""
+  app_root = tmp_path / "app"
+  app_root.mkdir(parents=True, exist_ok=True)
+  (app_root / "sub_agent.yaml").write_text("name: sub\n")
+  (app_root / "tool.yml").write_text("name: tool\n")
+
+  response = builder_test_client.get(
+      "/builder/app/app?file_path=sub_agent.yaml"
+  )
+  assert response.status_code == 200
+  assert response.text == "name: sub\n"
+
+  response = builder_test_client.get("/builder/app/app?file_path=tool.yml")
+  assert response.status_code == 200
+  assert response.text == "name: tool\n"
+
+
 def test_builder_endpoints_not_registered_without_web(
     mock_session_service,
     mock_artifact_service,