From 6bcfd00d2d06124e87cac6934f7c008535a3ce98 Mon Sep 17 00:00:00 2001
From: Nathan Brake <33383515+njbrake@users.noreply.github.com>
Date: Wed, 23 Apr 2025 15:07:24 -0400
Subject: [PATCH 1/2] Prevent MCP ClientSession hang

Per https://modelcontextprotocol.io/specification/draft/basic/lifecycle#timeouts

"Implementations SHOULD establish timeouts for all sent requests, to prevent hung connections and resource exhaustion. When the request has not received a success or error response within the timeout period, the sender SHOULD issue a cancellation notification for that request and stop waiting for a response.

SDKs and other middleware SHOULD allow these timeouts to be configured on a per-request basis."

I picked 5 seconds since that's the default for SSE
---
 src/google/adk/tools/mcp_tool/mcp_session_manager.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py
index a318b4b7..669a0982 100644
--- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py
+++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py
@@ -2,6 +2,7 @@
 import functools
 import sys
 from typing import Any, TextIO
+from datetime import timedelta
 import anyio
 from pydantic import BaseModel
 
@@ -133,6 +134,7 @@ async def create_session(self) -> ClientSession:
         connection_params=self.connection_params,
         exit_stack=self.exit_stack,
         errlog=self.errlog,
+
     )
 
   @classmethod
@@ -142,6 +144,7 @@ async def initialize_session(
       connection_params: StdioServerParameters | SseServerParams,
       exit_stack: AsyncExitStack,
       errlog: TextIO = sys.stderr,
+      timeout: int = 5,
   ) -> ClientSession:
     """Initializes an MCP client session.
 
@@ -150,6 +153,7 @@ async def initialize_session(
         exit_stack: AsyncExitStack to manage the session lifecycle.
         errlog: (Optional) TextIO stream for error logging. Use only for
           initializing a local stdio MCP session.
+        timeout: (Optional) Timeout for the session communication.
 
     Returns:
         ClientSession: The initialized MCP client session.
@@ -171,6 +175,7 @@ async def initialize_session(
       )
 
     transports = await exit_stack.enter_async_context(client)
-    session = await exit_stack.enter_async_context(ClientSession(*transports))
+
+    session = await exit_stack.enter_async_context(ClientSession(*transports, read_timeout_seconds=timedelta(seconds=timeout)))
     await session.initialize()
     return session

From e7996de7c35f04429a91dade8b21695161eb94a1 Mon Sep 17 00:00:00 2001
From: Nathan Brake <33383515+njbrake@users.noreply.github.com>
Date: Wed, 23 Apr 2025 15:07:24 -0400
Subject: [PATCH 2/2] Add timeout for MCP ClientSession

---
 src/google/adk/tools/mcp_tool/mcp_session_manager.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py
index a318b4b7..669a0982 100644
--- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py
+++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py
@@ -2,6 +2,7 @@
 import functools
 import sys
 from typing import Any, TextIO
+from datetime import timedelta
 import anyio
 from pydantic import BaseModel
 
@@ -133,6 +134,7 @@ async def create_session(self) -> ClientSession:
         connection_params=self.connection_params,
         exit_stack=self.exit_stack,
         errlog=self.errlog,
+
     )
 
   @classmethod
@@ -142,6 +144,7 @@ async def initialize_session(
       connection_params: StdioServerParameters | SseServerParams,
       exit_stack: AsyncExitStack,
       errlog: TextIO = sys.stderr,
+      timeout: int = 5,
   ) -> ClientSession:
     """Initializes an MCP client session.
 
@@ -150,6 +153,7 @@ async def initialize_session(
         exit_stack: AsyncExitStack to manage the session lifecycle.
         errlog: (Optional) TextIO stream for error logging. Use only for
           initializing a local stdio MCP session.
+        timeout: (Optional) Timeout for the session communication.
 
     Returns:
         ClientSession: The initialized MCP client session.
@@ -171,6 +175,7 @@ async def initialize_session(
       )
 
     transports = await exit_stack.enter_async_context(client)
-    session = await exit_stack.enter_async_context(ClientSession(*transports))
+
+    session = await exit_stack.enter_async_context(ClientSession(*transports, read_timeout_seconds=timedelta(seconds=timeout)))
     await session.initialize()
     return session