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
@@ -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."""