fix(mcp): recover pooled session after server-side termination

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

@@ -66,6 +66,7 @@ class AsyncAuthorizedSession:  # pylint: disable=g-bad-classes
 from mcp.client.stdio import stdio_client
 from mcp.client.streamable_http import create_mcp_http_client as _create_mcp_http_client
 from mcp.client.streamable_http import streamable_http_client
+from mcp.shared.exceptions import McpError
 from pydantic import BaseModel
 from pydantic import ConfigDict
 
@@ -487,6 +488,23 @@ async def wrapper(self, *args, **kwargs):
   return wrapper
 
 
+def _is_session_terminated_error(error: BaseException) -> bool:
+  """Whether `error` means the server no longer knows the session.
+
+  The streamable-HTTP client maps a 404 for a stored `mcp-session-id` (server
+  restart, idle-session eviction) to a JSON-RPC error with this message while
+  leaving the local read/write streams open and the background task alive, so
+  the pooled session passes every local health check yet fails every call.
+
+  Args:
+      error: The exception raised by a call on the session.
+
+  Returns:
+      True if the session should be dropped from the pool.
+  """
+  return isinstance(error, McpError) and 'Session terminated' in str(error)
+
+
 def _is_google_api_host(host: str | None) -> bool:
   """Returns whether host is a Google API endpoint."""
   if not host:
@@ -990,6 +1008,26 @@ def _end_session_use(self, headers: Optional[Dict[str, str]] = None) -> None:
     if session_key in self._sessions:
       self._session_last_used[session_key] = time.monotonic()
 
+  async def _invalidate_session(
+      self, headers: Optional[Dict[str, str]] = None
+  ) -> None:
+    """Drops the pooled session for these headers so the next call rebuilds it.
+
+    Needed when the server reports it no longer knows the session (see
+    `_is_session_terminated_error`): the local streams and background task
+    still look healthy, so `create_session` would keep returning the dead
+    session from the pool.
+
+    Args:
+        headers: Optional headers identifying the session, exactly as they
+          would be passed to ``create_session``.
+    """
+    session_key = self._session_key_for(headers)
+    async with self._session_lock:
+      if session_key in self._sessions:
+        _, exit_stack, stored_loop = self._sessions[session_key]
+        await self._cleanup_session(session_key, exit_stack, stored_loop)
+
   async def _cleanup_session(
       self,
       session_key: str,

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

@@ -29,6 +29,7 @@
 from google.genai.types import FunctionDeclaration
 from mcp import ClientSession
 from mcp.shared.exceptions import McpError
+from mcp.types import CallToolResult
 from mcp.types import Tool as McpBaseTool
 from opentelemetry import propagate
 from typing_extensions import override
@@ -56,6 +57,7 @@
 from ..tool_context import ToolContext
 from ..transfer_to_agent_tool import transfer_to_agent
 from .mcp_session_manager import _http_debug_var
+from .mcp_session_manager import _is_session_terminated_error
 from .mcp_session_manager import MCPSessionManager
 from .mcp_session_manager import retry_on_errors
 from .session_context import SessionContext
@@ -521,6 +523,59 @@ async def _run_async_impl(
     # Resolve progress callback (may be a factory that needs runtime context)
     resolved_callback = self._resolve_progress_callback(tool_context)
 
+    try:
+      response = await self._call_tool_on_session(
+          session, final_headers, args, resolved_callback, meta_trace_context
+      )
+    except McpError as e:
+      if not _is_session_terminated_error(e):
+        raise
+      # The server rejected the pooled session id (restart or idle eviction)
+      # before running the tool, so no side effect happened and one retry on
+      # a fresh session is safe. `_call_tool_on_session` already dropped the
+      # dead session from the pool, so `_create_session` builds a new one.
+      logger.info(
+          "MCP session was terminated server-side; retrying %s on a fresh"
+          " session.",
+          self._mcp_tool.name,
+      )
+      session = await self._create_session(headers=final_headers)
+      response = await self._call_tool_on_session(
+          session, final_headers, args, resolved_callback, meta_trace_context
+      )
+
+    result = response.model_dump(exclude_none=True, mode="json")
+
+    # Push UI widget to the event actions if the tool supports it.
+    if self.mcp_app_resource_uri:
+      tool_context.render_ui_widget(
+          UiWidget(
+              id=tool_context.function_call_id,
+              provider="mcp",
+              payload={
+                  "resource_uri": self.mcp_app_resource_uri,
+                  "tool": self._mcp_tool,
+                  "tool_args": args,
+              },
+          )
+      )
+    return result
+
+  async def _call_tool_on_session(
+      self,
+      session: ClientSession,
+      final_headers: dict[str, str] | None,
+      args: dict[str, Any],
+      resolved_callback: ProgressFnT | None,
+      meta_trace_context: dict[str, str] | None,
+  ) -> CallToolResult:
+    """Runs one tool call on `session`.
+
+    If the server reports the session as terminated (it restarted or evicted
+    the session id), the pooled session is dropped before the error
+    propagates: its local streams and background task still look healthy, so
+    without this every later call would keep reusing the dead session.
+    """
     call_coro = session.call_tool(
         self._mcp_tool.name,
         arguments=args,
@@ -552,33 +607,21 @@ async def _run_async_impl(
             headers=final_headers
         )
         if isinstance(session_context, SessionContext):
-          response = await session_context._run_guarded(call_coro)  # pylint: disable=protected-access
-        else:
-          response = await call_coro
+          return await session_context._run_guarded(call_coro)  # pylint: disable=protected-access
+        return await call_coro
       else:
         # Pre-fix behavior: await the call directly. This is what causes the
         # ~300s hang when the underlying transport crashes.
-        response = await call_coro
+        return await call_coro
+    except McpError as e:
+      if _is_session_terminated_error(e):
+        await self._mcp_session_manager._invalidate_session(  # pylint: disable=protected-access
+            headers=final_headers
+        )
+      raise
     finally:
       self._mcp_session_manager._end_session_use(final_headers)  # pylint: disable=protected-access
 
-    result = response.model_dump(exclude_none=True, mode="json")
-
-    # Push UI widget to the event actions if the tool supports it.
-    if self.mcp_app_resource_uri:
-      tool_context.render_ui_widget(
-          UiWidget(
-              id=tool_context.function_call_id,
-              provider="mcp",
-              payload={
-                  "resource_uri": self.mcp_app_resource_uri,
-                  "tool": self._mcp_tool,
-                  "tool_args": args,
-              },
-          )
-      )
-    return result
-
   def _detect_error_in_response(self, response: Any) -> str | None:
     """Telemetry hook: returns an error type if the response indicates an error."""
     # `response` is a dumped CallToolResult. MCP SDK 1.x names the field

tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

@@ -34,6 +34,7 @@
 from google.adk.tools.mcp_tool.mcp_session_manager import _DebugHttpxClientFactory
 from google.adk.tools.mcp_tool.mcp_session_manager import _GoogleAuthAsyncByteStream
 from google.adk.tools.mcp_tool.mcp_session_manager import _http_debug_var
+from google.adk.tools.mcp_tool.mcp_session_manager import _is_session_terminated_error
 from google.adk.tools.mcp_tool.mcp_session_manager import _RefreshableAsyncCredentials
 from google.adk.tools.mcp_tool.mcp_session_manager import _sanitize_url
 from google.adk.tools.mcp_tool.mcp_session_manager import _SESSION_IDLE_TTL_SECONDS
@@ -49,6 +50,8 @@
 from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
 import httpx
 from mcp import StdioServerParameters
+from mcp.shared.exceptions import McpError
+from mcp.types import ErrorData
 import pytest
 
 try:
@@ -1558,6 +1561,48 @@ async def mock_function(self):
   assert call_count == 1
 
 
+def test_is_session_terminated_error():
+  """Only the SDK's session-terminated error marks a session as dead."""
+  assert _is_session_terminated_error(
+      McpError(ErrorData(code=32600, message="Session terminated"))
+  )
+  assert not _is_session_terminated_error(
+      McpError(ErrorData(code=-32602, message="Invalid request parameters"))
+  )
+  assert not _is_session_terminated_error(ConnectionError("Session terminated"))
+
+
+@pytest.mark.asyncio
+async def test_invalidate_session_drops_pooled_session():
+  """Regression test for https://github.com/google/adk-python/issues/6822.
+
+  A server-terminated streamable-HTTP session keeps its local streams open
+  and its background task alive, so `create_session`'s health checks keep
+  returning it. `_invalidate_session` is the explicit path that drops it
+  from the pool so the next call builds a fresh session.
+  """
+  manager = MCPSessionManager(
+      StreamableHTTPConnectionParams(url="http://example.com/mcp")
+  )
+  session_key = manager._session_key_for(None)
+  exit_stack = MockAsyncExitStack()
+  manager._sessions[session_key] = (
+      MockClientSession(),
+      exit_stack,
+      asyncio.get_running_loop(),
+  )
+  manager._session_last_used[session_key] = time.monotonic()
+
+  await manager._invalidate_session()
+
+  assert session_key not in manager._sessions
+  assert session_key not in manager._session_last_used
+  exit_stack.aclose.assert_awaited_once()
+
+  # Invalidating an absent session is a harmless no-op.
+  await manager._invalidate_session()
+
+
 @pytest.mark.asyncio
 async def test_retry_on_errors_decorator_does_not_retry_when_task_is_cancelling():
   """Test the retry_on_errors decorator does not retry when cancelling."""

tests/unittests/tools/mcp_tool/test_mcp_tool.py

@@ -39,7 +39,9 @@
 from google.adk.tools.mcp_tool.mcp_tool import ProgressFnT
 from google.adk.tools.tool_context import ToolContext
 from google.genai.types import FunctionDeclaration
+from mcp.shared.exceptions import McpError
 from mcp.types import CallToolResult
+from mcp.types import ErrorData
 from mcp.types import TextContent
 from mcp.types import Tool as McpBaseTool
 import pytest
@@ -976,6 +978,104 @@ async def test_run_async_impl_retries_session_setup(self):
     assert self.mock_session_manager.create_session.await_count == 2
     self.mock_session.call_tool.assert_awaited_once()
 
+  @pytest.mark.asyncio
+  async def test_run_async_impl_recovers_from_terminated_session(self):
+    """Regression test for https://github.com/google/adk-python/issues/6822.
+
+    A server-side session termination (restart or idle eviction) leaves the
+    pooled session locally healthy but rejected by the server. The dead
+    session must be dropped from the pool and the call retried once on a
+    fresh session: the server refused the request before running the tool,
+    so the retry cannot duplicate a side effect.
+    """
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+    dead_session = AsyncMock()
+    dead_session.call_tool = AsyncMock(
+        side_effect=McpError(
+            ErrorData(code=32600, message="Session terminated")
+        )
+    )
+    fresh_session = AsyncMock()
+    response = Mock()
+    response.model_dump.return_value = {"result": "ok"}
+    fresh_session.call_tool = AsyncMock(return_value=response)
+    self.mock_session_manager.create_session = AsyncMock(
+        side_effect=[dead_session, fresh_session]
+    )
+    self.mock_session_manager._invalidate_session = AsyncMock()
+    tool_context = ToolContext(invocation_context=Mock())
+
+    result = await tool._run_async_impl(
+        args={"param1": "test_value"},
+        tool_context=tool_context,
+        credential=None,
+    )
+
+    assert result == {"result": "ok"}
+    self.mock_session_manager._invalidate_session.assert_awaited_once_with(
+        headers=None
+    )
+    assert self.mock_session_manager.create_session.await_count == 2
+    fresh_session.call_tool.assert_awaited_once()
+
+  @pytest.mark.asyncio
+  async def test_run_async_impl_does_not_retry_other_mcp_errors(self):
+    """An ordinary MCP protocol error keeps the pooled session and no retry."""
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+    self.mock_session.call_tool = AsyncMock(
+        side_effect=McpError(
+            ErrorData(code=-32602, message="Invalid request parameters")
+        )
+    )
+    self.mock_session_manager._invalidate_session = AsyncMock()
+    tool_context = ToolContext(invocation_context=Mock())
+
+    with pytest.raises(McpError, match="Invalid request parameters"):
+      await tool._run_async_impl(
+          args={"param1": "test_value"},
+          tool_context=tool_context,
+          credential=None,
+      )
+
+    self.mock_session_manager._invalidate_session.assert_not_awaited()
+    self.mock_session_manager.create_session.assert_awaited_once_with(
+        headers=None
+    )
+    self.mock_session.call_tool.assert_awaited_once()
+
+  @pytest.mark.asyncio
+  async def test_run_async_impl_terminated_session_twice_raises(self):
+    """A second terminated-session failure surfaces instead of looping."""
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+    self.mock_session.call_tool = AsyncMock(
+        side_effect=McpError(
+            ErrorData(code=32600, message="Session terminated")
+        )
+    )
+    self.mock_session_manager._invalidate_session = AsyncMock()
+    tool_context = ToolContext(invocation_context=Mock())
+
+    with pytest.raises(McpError, match="Session terminated"):
+      await tool._run_async_impl(
+          args={"param1": "test_value"},
+          tool_context=tool_context,
+          credential=None,
+      )
+
+    # Both attempts invalidated the dead pool entry; only one retry happened.
+    assert self.mock_session_manager._invalidate_session.await_count == 2
+    assert self.mock_session_manager.create_session.await_count == 2
+    assert self.mock_session.call_tool.await_count == 2
+
   @pytest.mark.asyncio
   async def test_get_headers_http_custom_scheme(self):
     """Test header generation for custom HTTP scheme."""