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 d95d48f2..7d9714aa 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/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 diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index ad420a3d..284aea41 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -34,7 +34,7 @@ from ...features import FeatureName from ...features import is_feature_enabled 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. @@ -195,7 +195,7 @@ class McpTool(BaseAuthenticatedTool): 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 diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index daa88f90..3768477e 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -37,7 +37,7 @@ from ..base_toolset import ToolPredicate 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 @@ class McpToolset(BaseToolset): 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, diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index 6c001ccf..b2d6b1cb 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -32,7 +32,7 @@ pytestmark = pytest.mark.skipif( # 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 @@ except ImportError as e: 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 @@ class TestMCPSessionManager: 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