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: 833487825
This commit is contained in:
Kathy Wu
2025-11-17 14:05:02 -08:00
committed by Copybara-Service
parent a5ac1d5e14
commit a48a1a9e88
3 changed files with 13 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/tools/mcp_tool/test_mcp_session_manager.py
+9 -48
View File
@@ -54,7 +54,6 @@ 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:
@@ -63,9 +62,6 @@ except ImportError:
self.command = command
self.args = args or []
class EmptyResult:
pass
class MockClientSession:
"""Mock ClientSession for testing."""
@@ -76,7 +72,6 @@ class MockClientSession:
self._read_stream._closed = False
self._write_stream._closed = False
self.initialize = AsyncMock()
self.send_ping = AsyncMock()
class MockAsyncExitStack:
@@ -211,52 +206,19 @@ class TestMCPSessionManager:
}
assert merged == expected
@pytest.mark.asyncio
async def test_is_session_disconnected_when_connected(self):
"""Test session disconnection detection when session is connected."""
def test_is_session_disconnected(self):
"""Test session disconnection detection."""
manager = MCPSessionManager(self.mock_stdio_connection_params)
session = MockClientSession()
session.send_ping.return_value = EmptyResult()
assert not await manager._is_session_disconnected(session)
session.send_ping.assert_called_once()
@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)
# Create mock session
session = MockClientSession()
session.send_ping.return_value = EmptyResult()
# Not disconnected
assert not manager._is_session_disconnected(session)
# Disconnected - read stream closed
session._read_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_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()
assert manager._is_session_disconnected(session)
@pytest.mark.asyncio
async def test_create_session_stdio_new(self):
@@ -309,7 +271,6 @@ 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()