Merge branch 'main' into fix-cli-robustness

Author: EaCognitive

.github/workflows/pre-commit.yml

@@ -0,0 +1,39 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: pre-commit
+
+on:
+  push:
+    branches: [main, v2]
+    paths:
+      - '**.py'
+      - '.pre-commit-config.yaml'
+      - 'pyproject.toml'
+  pull_request:
+    branches: [main, v2]
+    paths:
+      - '**.py'
+      - '.pre-commit-config.yaml'
+      - 'pyproject.toml'
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+
+      - name: Run pre-commit checks
+        uses: pre-commit/action@v3.0.1

pyproject.toml

@@ -90,6 +90,7 @@ dev = [
   "flit>=3.10.0",
   "isort>=6.0.0",
   "mypy>=1.15.0",
+  "pre-commit>=4.0.0",
   "pyink>=25.12.0",
   "pylint>=2.6.0",
   # go/keep-sorted end

src/google/adk/features/_feature_registry.py

@@ -41,6 +41,12 @@ class FeatureName(str, Enum):
   GOOGLE_CREDENTIALS_CONFIG = "GOOGLE_CREDENTIALS_CONFIG"
   GOOGLE_TOOL = "GOOGLE_TOOL"
   JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
+  # Private (leading underscore): not part of the public API surface.
+  # GE flips this on by setting the env var
+  # `ADK_ENABLE_MCP_GRACEFUL_ERROR_HANDLING=1`; nothing should import this
+  # enum member by name. Keeping it private avoids a backward-compat
+  # obligation for what is intended as a temporary, internal kill-switch.
+  _MCP_GRACEFUL_ERROR_HANDLING = "MCP_GRACEFUL_ERROR_HANDLING"
   PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING"
   PUBSUB_TOOL_CONFIG = "PUBSUB_TOOL_CONFIG"
   PUBSUB_TOOLSET = "PUBSUB_TOOLSET"
@@ -131,6 +137,9 @@ class FeatureConfig:
     FeatureName.JSON_SCHEMA_FOR_FUNC_DECL: FeatureConfig(
         FeatureStage.WIP, default_on=False
     ),
+    FeatureName._MCP_GRACEFUL_ERROR_HANDLING: FeatureConfig(
+        FeatureStage.EXPERIMENTAL, default_on=False
+    ),
     FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig(
         FeatureStage.EXPERIMENTAL, default_on=True
     ),

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -1949,6 +1949,7 @@ def __init__(
       table_id: Optional[str] = None,
       config: Optional[BigQueryLoggerConfig] = None,
       location: str = "US",
+      credentials: Optional[google.auth.credentials.Credentials] = None,
       **kwargs,
   ) -> None:
     """Initializes the instance.
@@ -1959,6 +1960,8 @@ def __init__(
         table_id: BigQuery table ID (optional, overrides config).
         config: BigQueryLoggerConfig (optional).
         location: BigQuery location (default: "US").
+        credentials: Google Auth credentials (optional). If None, uses
+          Application Default Credentials.
         **kwargs: Additional configuration parameters for BigQueryLoggerConfig.
     """
     super().__init__(name="bigquery_agent_analytics")
@@ -1985,6 +1988,7 @@ def __init__(
     self._startup_error: Optional[Exception] = None
     self._is_shutting_down = False
     self._setup_lock = None
+    self._credentials = credentials
     self.client = None
     self._loop_state_by_loop: dict[asyncio.AbstractEventLoop, _LoopState] = {}
     self._write_stream_name = None  # Resolved stream name
@@ -2097,15 +2101,16 @@ async def _get_loop_state(self) -> _LoopState:
     # grpc.aio clients are loop-bound, so we create one per event loop.
 
     def get_credentials():
-      creds, project_id = google.auth.default(
+      creds, _ = google.auth.default(
           scopes=["https://www.googleapis.com/auth/cloud-platform"]
       )
-      return creds, project_id
+      return creds
 
-    creds, project_id = await loop.run_in_executor(
-        self._executor, get_credentials
-    )
-    quota_project_id = getattr(creds, "quota_project_id", None)
+    if self._credentials is None:
+      self._credentials = await loop.run_in_executor(
+          self._executor, get_credentials
+      )
+    quota_project_id = getattr(self._credentials, "quota_project_id", None)
     options = (
         client_options.ClientOptions(quota_project_id=quota_project_id)
         if quota_project_id
@@ -2119,7 +2124,7 @@ def get_credentials():
     client_info = gapic_client_info.ClientInfo(user_agent=" ".join(user_agents))
 
     write_client = BigQueryWriteAsyncClient(
-        credentials=creds,
+        credentials=self._credentials,
         client_info=client_info,
         client_options=options,
     )
@@ -2173,7 +2178,9 @@ async def _lazy_setup(self, **kwargs) -> None:
       self.client = await loop.run_in_executor(
           self._executor,
           lambda: bigquery.Client(
-              project=self.project_id, location=self.location
+              project=self.project_id,
+              location=self.location,
+              credentials=self._credentials,
           ),
       )
 
@@ -2193,7 +2200,9 @@ async def _lazy_setup(self, **kwargs) -> None:
             self.project_id,
             self.config.gcs_bucket_name,
             self._executor,
-            storage_client=kwargs.get("storage_client"),
+            storage_client=storage.Client(
+                project=self.project_id, credentials=self._credentials
+            ),
         )
 
       self.parser = HybridContentParser(

src/google/adk/sessions/vertex_ai_session_service.py

@@ -48,6 +48,12 @@
 _USAGE_METADATA_CUSTOM_METADATA_KEY = '_usage_metadata'
 
 
+def _quote_filter_literal(value: str) -> str:
+  """Quotes filter values so embedded metacharacters stay inside the literal."""
+  escaped_value = value.replace('\\', '\\\\').replace('"', '\\"')
+  return f'"{escaped_value}"'
+
+
 def _set_internal_custom_metadata(
     metadata_dict: dict[str, Any], *, key: str, value: dict[str, Any]
 ) -> None:
@@ -228,7 +234,7 @@ async def list_sessions(
       sessions = []
       config = {}
       if user_id is not None:
-        config['filter'] = f'user_id="{user_id}"'
+        config['filter'] = f'user_id={_quote_filter_literal(user_id)}'
       sessions_iterator = await api_client.agent_engines.sessions.list(
           name=f'reasoningEngines/{reasoning_engine_id}',
           config=config,

src/google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -44,6 +44,8 @@
 from pydantic import BaseModel
 from pydantic import ConfigDict
 
+from ...features import FeatureName
+from ...features import is_feature_enabled
 from .session_context import SessionContext
 
 logger = logging.getLogger('google_adk.' + __name__)
@@ -237,11 +239,18 @@ def __init__(
       self._connection_params = connection_params
     self._errlog = errlog
 
-    # Session pool: maps session keys to (session, exit_stack, loop) tuples
+    # Session pool: maps session keys to (session, exit_stack, loop) tuples.
+    # Kept as a tuple for backward-compatibility with downstream tests
+    # that construct or unpack entries directly.
     self._sessions: Dict[
         str, tuple[ClientSession, AsyncExitStack, asyncio.AbstractEventLoop]
     ] = {}
 
+    # Sibling pool: maps session keys to their SessionContext. Stored
+    # separately from `_sessions` so the tuple shape above stays stable.
+    # Used by McpTool to access `_run_guarded` for transport-crash detection.
+    self._session_contexts: Dict[str, SessionContext] = {}
+
     # Map of event loops to their respective locks to prevent race conditions
     # across different event loops in session creation.
     self._session_lock_map: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
@@ -323,6 +332,26 @@ def _is_session_disconnected(self, session: ClientSession) -> bool:
     """
     return session._read_stream._closed or session._write_stream._closed
 
+  def _get_session_context(
+      self, headers: Optional[Dict[str, str]] = None
+  ) -> Optional[SessionContext]:
+    """Returns the SessionContext for the session matching the given headers.
+
+    Note: This method reads from the session-context pool without acquiring
+    ``_session_lock``. This is safe because it is called immediately after
+    ``create_session()`` (which populates the entry under the lock) within
+    the same task, and dict reads are atomic in CPython.
+
+    Args:
+        headers: Optional headers used to identify the session.
+
+    Returns:
+        The SessionContext if a matching session exists, None otherwise.
+    """
+    merged_headers = self._merge_headers(headers)
+    session_key = self._generate_session_key(merged_headers)
+    return self._session_contexts.get(session_key)
+
   async def _cleanup_session(
       self,
       session_key: str,
@@ -378,6 +407,10 @@ def cleanup_done(f: asyncio.Future):
     finally:
       if session_key in self._sessions:
         del self._sessions[session_key]
+      # Also drop the SessionContext reference so we don't leak the
+      # SessionContext after its underlying session is gone.
+      if session_key in self._session_contexts:
+        del self._session_contexts[session_key]
 
   def _create_client(self, merged_headers: Optional[Dict[str, str]] = None):
     """Creates an MCP client based on the connection parameters.
@@ -453,15 +486,30 @@ async def create_session(
       if session_key in self._sessions:
         session, exit_stack, stored_loop = self._sessions[session_key]
 
-        # Check if the existing session is still connected and bound to the current loop
+        # Check if the existing session is still connected and bound to
+        # the current loop. When the feature flag is on, we ALSO check the
+        # SessionContext's background task: a crashed transport can leave
+        # the session's read/write streams open even though the underlying
+        # task has already died (e.g. after a 4xx/5xx HTTP response).
+        # Without that extra check, callers would reuse a dead session and
+        # hang on the next call. The check is gated because it triggers
+        # session re-creation in some test mocks where `_task` looks
+        # "not alive" but the streams are otherwise reusable.
         current_loop = asyncio.get_running_loop()
-        if stored_loop is current_loop and not self._is_session_disconnected(
-            session
+        if is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING):  # pylint: disable=protected-access
+          ctx = self._session_contexts.get(session_key)
+          ctx_alive = ctx is None or ctx._is_task_alive  # pylint: disable=protected-access
+        else:
+          ctx_alive = True  # Pre-fix: do not consult task aliveness
+        if (
+            stored_loop is current_loop
+            and not self._is_session_disconnected(session)
+            and ctx_alive
         ):
           # Session is still good, return it
           return session
         else:
-          # Session is disconnected or from a different loop, clean it up
+          # Session is disconnected, dead, or from a different loop; clean up.
           logger.info(
               'Cleaning up session (disconnected or different loop): %s',
               session_key,
@@ -485,26 +533,32 @@ async def create_session(
         client = self._create_client(merged_headers)
         is_stdio = isinstance(self._connection_params, StdioConnectionParams)
 
+        session_context = SessionContext(
+            client=client,
+            timeout=timeout_in_seconds,
+            sse_read_timeout=sse_read_timeout_in_seconds,
+            is_stdio=is_stdio,
+            sampling_callback=self._sampling_callback,
+            sampling_capabilities=self._sampling_capabilities,
+        )
+
         session = await asyncio.wait_for(
-            exit_stack.enter_async_context(
-                SessionContext(
-                    client=client,
-                    timeout=timeout_in_seconds,
-                    sse_read_timeout=sse_read_timeout_in_seconds,
-                    is_stdio=is_stdio,
-                    sampling_callback=self._sampling_callback,
-                    sampling_capabilities=self._sampling_capabilities,
-                )
-            ),
+            exit_stack.enter_async_context(session_context),
             timeout=timeout_in_seconds,
         )
 
-        # Store session, exit stack, and loop in the pool
+        # Store session, exit stack, and loop in the pool. The pool storage
+        # remains a tuple for backward-compatibility with downstream tests
+        # that construct or unpack entries directly.
         self._sessions[session_key] = (
             session,
             exit_stack,
             asyncio.get_running_loop(),
         )
+        # Track the SessionContext in a sibling dict so McpTool can call
+        # `_run_guarded` on it. Stored separately to avoid changing the
+        # shape of `_sessions` (which is a public-ish internal surface).
+        self._session_contexts[session_key] = session_context
         logger.debug('Created new session: %s', session_key)
         return session
 
@@ -524,6 +578,7 @@ def __getstate__(self):
     state = self.__dict__.copy()
     # Remove unpicklable entries or those that shouldn't persist across pickle
     state['_sessions'] = {}
+    state['_session_contexts'] = {}
     state['_session_lock_map'] = {}
 
     # Locks and file-like objects cannot be pickled
@@ -537,6 +592,7 @@ def __setstate__(self, state):
     self.__dict__.update(state)
     # Re-initialize members that were not pickled
     self._sessions = {}
+    self._session_contexts = {}
     self._session_lock_map = {}
     self._lock_map_lock = threading.Lock()
     # If _errlog was removed during pickling, default to sys.stderr