mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
99ca6aa6e6
commit
ee8106be77
@@ -339,38 +339,50 @@ class MCPSessionManager:
|
||||
|
||||
# Create a new session (either first time or replacing disconnected one)
|
||||
exit_stack = AsyncExitStack()
|
||||
timeout_in_seconds = (
|
||||
self._connection_params.timeout
|
||||
if hasattr(self._connection_params, 'timeout')
|
||||
else None
|
||||
)
|
||||
|
||||
try:
|
||||
client = self._create_client(merged_headers)
|
||||
|
||||
transports = await exit_stack.enter_async_context(client)
|
||||
# The streamable http client returns a GetSessionCallback in addition to the read/write MemoryObjectStreams
|
||||
# needed to build the ClientSession, we limit then to the two first values to be compatible with all clients.
|
||||
transports = await asyncio.wait_for(
|
||||
exit_stack.enter_async_context(client),
|
||||
timeout=timeout_in_seconds,
|
||||
)
|
||||
# The streamable http client returns a GetSessionCallback in addition to the
|
||||
# read/write MemoryObjectStreams needed to build the ClientSession, we limit
|
||||
# then to the two first values to be compatible with all clients.
|
||||
if isinstance(self._connection_params, StdioConnectionParams):
|
||||
session = await exit_stack.enter_async_context(
|
||||
ClientSession(
|
||||
*transports[:2],
|
||||
read_timeout_seconds=timedelta(
|
||||
seconds=self._connection_params.timeout
|
||||
),
|
||||
read_timeout_seconds=timedelta(seconds=timeout_in_seconds),
|
||||
)
|
||||
)
|
||||
else:
|
||||
session = await exit_stack.enter_async_context(
|
||||
ClientSession(*transports[:2])
|
||||
)
|
||||
await session.initialize()
|
||||
await asyncio.wait_for(session.initialize(), timeout=timeout_in_seconds)
|
||||
|
||||
# Store session and exit stack in the pool
|
||||
self._sessions[session_key] = (session, exit_stack)
|
||||
logger.debug('Created new session: %s', session_key)
|
||||
return session
|
||||
|
||||
except Exception:
|
||||
except Exception as e:
|
||||
# If session creation fails, clean up the exit stack
|
||||
if exit_stack:
|
||||
await exit_stack.aclose()
|
||||
raise
|
||||
try:
|
||||
await exit_stack.aclose()
|
||||
except Exception as exit_stack_error:
|
||||
logger.warning(
|
||||
'Error during session creation cleanup: %s', exit_stack_error
|
||||
)
|
||||
raise ConnectionError(f'Failed to create MCP session: {e}') from e
|
||||
|
||||
async def close(self):
|
||||
"""Closes all sessions and cleans up resources."""
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user