diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 1a2b188180..6fcef3d153 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/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, diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index e92f144cc1..aad70dd9c6 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/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 diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index 32527ecd1b..ea4a729600 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/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.""" diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 4e8564dce1..de4dd953d5 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/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."""