Merge branch 'main' into fix/anthropic-thinking-streaming-continuity

Author: gurjot-05

src/google/adk/agents/base_agent.py

@@ -16,7 +16,6 @@
 
 import inspect
 import logging
-import sys
 from typing import Any
 from typing import AsyncGenerator
 from typing import Awaitable
@@ -286,9 +285,7 @@ async def run_async(
       Event: the events generated by the agent.
     """
 
-    cm = tracer.start_as_current_span(f'invoke_agent {self.name}')
-    span = cm.__enter__()
-    try:
+    with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
       ctx = self._create_invocation_context(parent_context)
       tracing.trace_agent_invocation(span, self, ctx)
       if event := await self._handle_before_agent_callback(ctx):
@@ -305,23 +302,6 @@ async def run_async(
 
       if event := await self._handle_after_agent_callback(ctx):
         yield event
-    except BaseException:
-      try:
-        cm.__exit__(*sys.exc_info())
-      except ValueError:
-        logger.warning(
-            'Failed to detach context during generator cleanup, likely due to'
-            ' cancellation.'
-        )
-      raise
-    else:
-      try:
-        cm.__exit__(None, None, None)
-      except ValueError:
-        logger.warning(
-            'Failed to detach context during generator cleanup, likely due to'
-            ' cancellation.'
-        )
 
   @final
   async def run_live(

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/flows/llm_flows/base_llm_flow.py

@@ -18,7 +18,6 @@
 import asyncio
 import inspect
 import logging
-import sys
 from typing import AsyncGenerator
 from typing import Optional
 from typing import TYPE_CHECKING
@@ -1169,9 +1168,7 @@ async def _call_llm_async(
   ) -> AsyncGenerator[LlmResponse, None]:
 
     async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
-      cm = tracer.start_as_current_span('call_llm')
-      span = cm.__enter__()
-      try:
+      with tracer.start_as_current_span('call_llm') as span:
         # Runs before_model_callback inside the call_llm span so
         # plugins observe the same span as after/error callbacks.
         if response := await self._handle_before_model_callback(
@@ -1264,23 +1261,6 @@ async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]:
                   llm_response = altered
 
               yield llm_response
-      except BaseException:
-        try:
-          cm.__exit__(*sys.exc_info())
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
-        raise
-      else:
-        try:
-          cm.__exit__(None, None, None)
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
 
     async with Aclosing(_call_llm_with_tracing()) as agen:
       async for event in agen:

src/google/adk/runners.py

@@ -543,9 +543,7 @@ async def _run_with_trace(
         new_message: Optional[types.Content] = None,
         invocation_id: Optional[str] = None,
     ) -> AsyncGenerator[Event, None]:
-      cm = tracer.start_as_current_span('invocation')
-      span = cm.__enter__()
-      try:
+      with tracer.start_as_current_span('invocation'):
         session = await self._get_or_create_session(
             user_id=user_id,
             session_id=session_id,
@@ -629,23 +627,6 @@ async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
               self.session_service,
               skip_token_compaction=invocation_context.token_compaction_checked,
           )
-      except BaseException:
-        try:
-          cm.__exit__(*sys.exc_info())
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
-        raise
-      else:
-        try:
-          cm.__exit__(None, None, None)
-        except ValueError:
-          logger.warning(
-              'Failed to detach context during generator cleanup, likely due to'
-              ' cancellation.'
-          )
 
     async with Aclosing(_run_with_trace(new_message, invocation_id)) as agen:
       async for event in agen:

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