fix: Fix error handling when MCP server is unreachable

Currently ADK web hangs and logs "AGSI callable returned without completing response" when the server is unreachable. To fix, set timeouts for connecting to server.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 828648928
This commit is contained in:
Kathy Wu
2025-11-05 15:15:58 -08:00
committed by Copybara-Service
parent 99ca6aa6e6
commit ee8106be77
2 changed files with 65 additions and 10 deletions
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
from datetime import timedelta
import hashlib
from io import StringIO
import json
@@ -279,6 +281,47 @@ class TestMCPSessionManager:
# Should not create new session
existing_session.initialize.assert_not_called()
@pytest.mark.asyncio
@patch("google.adk.tools.mcp_tool.mcp_session_manager.stdio_client")
@patch("google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack")
@patch("google.adk.tools.mcp_tool.mcp_session_manager.ClientSession")
async def test_create_session_timeout(
self, mock_session_class, mock_exit_stack_class, mock_stdio
):
"""Test session creation timeout."""
manager = MCPSessionManager(self.mock_stdio_connection_params)
mock_session = MockClientSession()
mock_exit_stack = MockAsyncExitStack()
mock_exit_stack_class.return_value = mock_exit_stack
mock_stdio.return_value = AsyncMock()
mock_exit_stack.enter_async_context.side_effect = [
("read", "write"), # First call returns transports
mock_session, # Second call returns session
]
mock_session_class.return_value = mock_session
# Simulate timeout during session initialization
mock_session.initialize.side_effect = asyncio.TimeoutError("Test timeout")
# Expect ConnectionError due to timeout
with pytest.raises(ConnectionError, match="Failed to create MCP session"):
await manager.create_session()
# Verify ClientSession called with timeout
mock_session_class.assert_called_with(
"read",
"write",
read_timeout_seconds=timedelta(
seconds=manager._connection_params.timeout
),
)
# Verify session was not added to pool
assert not manager._sessions
# Verify cleanup was called
mock_exit_stack.aclose.assert_called_once()
@pytest.mark.asyncio
async def test_close_success(self):
"""Test successful cleanup of all sessions."""