fix: Update the retry_on_closed_resource decorator to retry on all errors

Retrying only on closed_resource error is not enough to be reliable for production environments due to the other network errors that may occur -- remote protocol error, read timeout, etc. We will update this to retry on all errors. Since it is only a one-time retry, it should not affect latency significantly. Fixes #2561.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 831123085

Author: wukath

src/google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -108,10 +108,10 @@ class StreamableHTTPConnectionParams(BaseModel):
   terminate_on_close: bool = True
 
 
-def retry_on_closed_resource(func):
-  """Decorator to automatically retry action when MCP session is closed.
+def retry_on_errors(func):
+  """Decorator to automatically retry action when MCP session errors occur.
 
-  When MCP session was closed, the decorator will automatically retry the
+  When MCP session errors occur, the decorator will automatically retry the
   action once. The create_session method will handle creating a new session
   if the old one was disconnected.
 
@@ -126,11 +126,11 @@ def retry_on_closed_resource(func):
   async def wrapper(self, *args, **kwargs):
     try:
       return await func(self, *args, **kwargs)
-    except (anyio.ClosedResourceError, anyio.BrokenResourceError):
-      # If the session connection is closed or unusable, we will retry the
-      # function to reconnect to the server. create_session will handle
-      # detecting and replacing disconnected sessions.
-      logger.info('Retrying %s due to closed resource', func.__name__)
+    except Exception as e:
+      # If an error is thrown, we will retry the function to reconnect to the
+      # server. create_session will handle detecting and replacing disconnected
+      # sessions.
+      logger.info('Retrying %s due to error: %s', func.__name__, e)
       return await func(self, *args, **kwargs)
 
   return wrapper

src/google/adk/tools/mcp_tool/mcp_tool.py

@@ -32,7 +32,7 @@
 from ...agents.readonly_context import ReadonlyContext
 from .._gemini_schema_util import _to_gemini_schema
 from .mcp_session_manager import MCPSessionManager
-from .mcp_session_manager import retry_on_closed_resource
+from .mcp_session_manager import retry_on_errors
 
 # Attempt to import MCP Tool from the MCP library, and hints user to upgrade
 # their Python version to 3.10 if it fails.
@@ -184,7 +184,7 @@ async def run_async(
         return {"error": "This tool call is rejected."}
     return await super().run_async(args=args, tool_context=tool_context)
 
-  @retry_on_closed_resource
+  @retry_on_errors
   @override
   async def _run_async_impl(
       self, *, args, tool_context: ToolContext, credential: AuthCredential

src/google/adk/tools/mcp_tool/mcp_toolset.py

@@ -37,7 +37,7 @@
 from ..tool_configs import BaseToolConfig
 from ..tool_configs import ToolArgsConfig
 from .mcp_session_manager import MCPSessionManager
-from .mcp_session_manager import retry_on_closed_resource
+from .mcp_session_manager import retry_on_errors
 from .mcp_session_manager import SseConnectionParams
 from .mcp_session_manager import StdioConnectionParams
 from .mcp_session_manager import StreamableHTTPConnectionParams
@@ -155,7 +155,7 @@ def __init__(
     self._auth_credential = auth_credential
     self._require_confirmation = require_confirmation
 
-  @retry_on_closed_resource
+  @retry_on_errors
   async def get_tools(
       self,
       readonly_context: Optional[ReadonlyContext] = None,

tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

@@ -32,7 +32,7 @@
 # Import dependencies with version checking
 try:
   from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
-  from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_closed_resource
+  from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_errors
   from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
   from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
   from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
@@ -44,7 +44,7 @@ class DummyClass:
       pass
 
     MCPSessionManager = DummyClass
-    retry_on_closed_resource = lambda x: x
+    retry_on_errors = lambda x: x
     SseConnectionParams = DummyClass
     StdioConnectionParams = DummyClass
     StreamableHTTPConnectionParams = DummyClass
@@ -375,12 +375,12 @@ async def test_close_with_errors(self):
     assert "Close error 1" in error_output
 
 
-def test_retry_on_closed_resource_decorator():
-  """Test the retry_on_closed_resource decorator."""
+def test_retry_on_errors_decorator():
+  """Test the retry_on_errors decorator."""
 
   call_count = 0
 
-  @retry_on_closed_resource
+  @retry_on_errors
   async def mock_function(self):
     nonlocal call_count
     call_count += 1