encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: Improve logic for checking if a MCP session is disconnected

Browse Source 
Currently logic to check for a disconnected session only checks for certain headers but doesn't detect all cases, leading to situations where it tries to connect to a session that is down. This adds logic so that we ping the server to check if it is disconnected. Fixes https://github.com/google/adk-python/issues/3321.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 832460068
This commit is contained in:
Kathy Wu
2025-11-14 14:33:39 -08:00
committed by Copybara-Service
parent 29fea7ec1f
commit a754c96d3c
3 changed files with 73 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/tools/mcp_tool/test_mcp_session_manager.py
+48 -9
View File
@@ -54,6 +54,7 @@ except ImportError as e:
# Import real MCP classes
try:
from mcp import StdioServerParameters
from mcp.types import EmptyResult
except ImportError:
# Create a mock if MCP is not available
class StdioServerParameters:
@@ -62,6 +63,9 @@ except ImportError:
self.command = command
self.args = args or []
class EmptyResult:
pass
class MockClientSession:
"""Mock ClientSession for testing."""
@@ -72,6 +76,7 @@ class MockClientSession:
self._read_stream._closed = False
self._write_stream._closed = False
self.initialize = AsyncMock()
self.send_ping = AsyncMock()
class MockAsyncExitStack:
@@ -206,19 +211,52 @@ class TestMCPSessionManager:
}
assert merged == expected
def test_is_session_disconnected(self):
"""Test session disconnection detection."""
@pytest.mark.asyncio
async def test_is_session_disconnected_when_connected(self):
"""Test session disconnection detection when session is connected."""
manager = MCPSessionManager(self.mock_stdio_connection_params)
# Create mock session
session = MockClientSession()
session.send_ping.return_value = EmptyResult()
assert not await manager._is_session_disconnected(session)
session.send_ping.assert_called_once()
# Not disconnected
assert not manager._is_session_disconnected(session)
# Disconnected - read stream closed
@pytest.mark.asyncio
async def test_is_session_disconnected_read_stream_closed(self):
"""Test session disconnection detection when read stream is closed."""
manager = MCPSessionManager(self.mock_stdio_connection_params)
session = MockClientSession()
session.send_ping.return_value = EmptyResult()
session._read_stream._closed = True
assert manager._is_session_disconnected(session)
assert await manager._is_session_disconnected(session)
session.send_ping.assert_not_called()
@pytest.mark.asyncio
async def test_is_session_disconnected_write_stream_closed(self):
"""Test session disconnection detection when write stream is closed."""
manager = MCPSessionManager(self.mock_stdio_connection_params)
session = MockClientSession()
session.send_ping.return_value = EmptyResult()
session._write_stream._closed = True
assert await manager._is_session_disconnected(session)
session.send_ping.assert_not_called()
@pytest.mark.asyncio
async def test_is_session_disconnected_ping_fails(self):
"""Test session disconnection detection when ping fails."""
manager = MCPSessionManager(self.mock_stdio_connection_params)
session = MockClientSession()
session.send_ping.side_effect = Exception("Ping failed")
assert await manager._is_session_disconnected(session)
session.send_ping.assert_called_once()
@pytest.mark.asyncio
async def test_is_session_disconnected_ping_returns_wrong_result(self):
"""Test session disconnection detection when ping returns wrong result."""
manager = MCPSessionManager(self.mock_stdio_connection_params)
session = MockClientSession()
session.send_ping.return_value = "Wrong result"
assert await manager._is_session_disconnected(session)
session.send_ping.assert_called_once()
@pytest.mark.asyncio
async def test_create_session_stdio_new(self):
@@ -271,6 +309,7 @@ class TestMCPSessionManager:
# Session is connected
existing_session._read_stream._closed = False
existing_session._write_stream._closed = False
existing_session.send_ping.return_value = EmptyResult()
session = await manager.create_session()