fix: ADK 1.20 breaking change. Add retry_on_errors decorator for MCP session retries and fix (#348)

Introduced a new retry_on_errors decorator in utils.py to handle automatic retries on MCP session errors. Updated mcp_tool.py and mcp_toolset.py to use this decorator instead of retry_on_closed_resource, improving error handling and session recovery.
fix: ADK 1.20 breaking change on method  of sessionmanager

Author: loveyana

veadk/integrations/ve_identity/mcp_tool.py

@@ -24,14 +24,13 @@
 from google.adk.tools.tool_context import ToolContext
 from google.adk.tools._gemini_schema_util import _to_gemini_schema
 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 veadk.integrations.ve_identity.auth_config import VeIdentityAuthConfig
 from veadk.integrations.ve_identity.auth_mixins import (
     VeIdentityAuthMixin,
     AuthRequiredException,
 )
-from veadk.integrations.ve_identity.utils import generate_headers
+from veadk.integrations.ve_identity.utils import generate_headers, retry_on_errors
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -157,7 +156,7 @@ async def _execute_with_credential(
             args=args, tool_context=tool_context, credential=credential
         )
 
-    @retry_on_closed_resource
+    @retry_on_errors
     async def _run_async_impl(
         self, *, args, tool_context: ToolContext, credential: AuthCredential
     ):

veadk/integrations/ve_identity/mcp_toolset.py

@@ -29,7 +29,6 @@
 from mcp.types import ListToolsResult
 
 from google.adk.tools.mcp_tool.mcp_session_manager import (
-    retry_on_closed_resource,
     SseConnectionParams,
     StdioConnectionParams,
     StreamableHTTPConnectionParams,
@@ -43,7 +42,7 @@
 from veadk.integrations.ve_identity.auth_config import VeIdentityAuthConfig
 from veadk.integrations.ve_identity.auth_mixins import VeIdentityAuthMixin
 from veadk.integrations.ve_identity.mcp_tool import VeIdentityMcpTool
-from veadk.integrations.ve_identity.utils import generate_headers
+from veadk.integrations.ve_identity.utils import generate_headers, retry_on_errors
 
 logger = logging.getLogger(__name__)
 
@@ -158,7 +157,7 @@ def __init__(
             errlog=errlog,
         )
 
-    @retry_on_closed_resource
+    @retry_on_errors
     @override
     async def get_tools(
         self,

veadk/integrations/ve_identity/utils.py

@@ -16,12 +16,17 @@
 
 from __future__ import annotations
 import base64
+import functools
 from typing import Optional
 
 from google.adk.events import Event
 from google.adk.auth import AuthConfig
 from google.adk.auth.auth_credential import AuthCredential
 
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
 
 def is_pending_auth_event(event: Event) -> bool:
     """Check if an ADK event represents a pending authentication request.
@@ -149,3 +154,31 @@ def generate_headers(credential: AuthCredential) -> Optional[dict[str, str]]:
             pass
 
     return headers
+
+
+def retry_on_errors(func):
+    """Decorator to automatically retry action when MCP session errors occur.
+
+    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.
+
+    Args:
+        func: The function to decorate.
+
+    Returns:
+        The decorated function.
+    """
+
+    @functools.wraps(func)  # Preserves original function metadata
+    async def wrapper(self, *args, **kwargs):
+        try:
+            return await func(self, *args, **kwargs)
+        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