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 628c1b52..029f30f2 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -796,13 +796,30 @@ class MCPSessionManager: def _is_session_disconnected(self, session: ClientSession) -> bool: """Checks if a session is disconnected or closed. + Reads two attributes ADK does not own: the SDK holds the transport streams + on the session privately, and each stream reports its own closed flag. A + session that lacks either one reads as connected rather than raising, + because a release is free to restructure both away and this probe is not + the only thing standing between a dead session and a caller. + + `create_session` pairs this with `SessionContext._is_task_alive`, which + ADK owns and which catches strictly more: a crashed transport can leave + the streams open while the task behind them is already dead. That pairing + runs under `_MCP_GRACEFUL_ERROR_HANDLING`, which is on by default. The + kill switch drops it and leaves this probe on its own. + Args: session: The ClientSession to check. Returns: - True if the session is disconnected, False otherwise. + True if the session is known to be disconnected, False otherwise. """ - return session._read_stream._closed or session._write_stream._closed + read_stream = getattr(session, '_read_stream', None) + write_stream = getattr(session, '_write_stream', None) + return bool( + getattr(read_stream, '_closed', False) + or getattr(write_stream, '_closed', False) + ) def _get_session_context( self, headers: Optional[Dict[str, str]] = None 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 dd5dddf0..cdb97dd1 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -390,6 +390,47 @@ class TestMCPSessionManager: session._read_stream._closed = True assert manager._is_session_disconnected(session) + def test_is_session_disconnected_write_stream_closed(self): + """The write stream closing counts too, not just the read stream.""" + manager = MCPSessionManager(self.mock_stdio_connection_params) + + session = MockClientSession() + session._write_stream._closed = True + assert manager._is_session_disconnected(session) + + def test_is_session_disconnected_without_streams(self): + """A session that holds no streams reads as connected, and does not raise. + + Both attributes are private to the SDK. A release is free to move the + streams off `ClientSession`, and this must degrade to the + `SessionContext` task check rather than take down every tool call with an + `AttributeError`. + + The stand-in is a bare class on purpose: a `Mock` would answer to + `_read_stream` and pass this vacuously. + """ + + class SessionWithoutStreams: + pass + + manager = MCPSessionManager(self.mock_stdio_connection_params) + assert not manager._is_session_disconnected(SessionWithoutStreams()) + + def test_is_session_disconnected_with_streams_that_have_no_flag(self): + """A stream that stops reporting a closed flag reads as connected too.""" + + class StreamWithoutFlag: + pass + + class SessionWithBareStreams: + + def __init__(self): + self._read_stream = StreamWithoutFlag() + self._write_stream = StreamWithoutFlag() + + manager = MCPSessionManager(self.mock_stdio_connection_params) + assert not manager._is_session_disconnected(SessionWithBareStreams()) + @pytest.mark.asyncio async def test_create_session_stdio_new(self): """Test creating a new stdio session."""