feat: backport MCP client session reconnection (backport CF #3639)

Signed-off-by: Frederico Araujo <frederico.araujo@ibm.com>

Author: araujof

cpex/framework/external/mcp/client.py

@@ -21,7 +21,7 @@
 # Third-Party
 import httpx
 import orjson
-from mcp import ClientSession, StdioServerParameters
+from mcp import ClientSession, McpError, StdioServerParameters
 from mcp.client.stdio import stdio_client
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.types import TextContent
@@ -85,6 +85,8 @@ def __init__(self, config: PluginConfig) -> None:
         self._get_session_id: Optional[Callable[[], str | None]] = None
         self._session_id: Optional[str] = None
         self._http_client_factory: Optional[Callable[..., httpx.AsyncClient]] = None
+        self._reconnect_attempts: int = 3
+        self._reconnect_delay: float = 0.1
 
     async def initialize(self) -> None:
         """Initialize the plugin's connection to the MCP server.
@@ -99,6 +101,10 @@ async def initialize(self) -> None:
                     message="The mcp section must be defined for external plugin", plugin_name=self.name
                 )
             )
+
+        self._reconnect_attempts = self._config.mcp.reconnect_attempts
+        self._reconnect_delay = self._config.mcp.reconnect_delay
+
         if self._config.mcp.proto == TransportType.STDIO:
             if not (self._config.mcp.script or self._config.mcp.cmd):
                 raise PluginError(
@@ -405,6 +411,74 @@ def _tls_httpx_client_factory(
                 logger.info("Retrying in %ss...", delay)
                 await asyncio.sleep(delay)
 
+    async def _cleanup_session(self) -> None:
+        """Reset session state without a full shutdown (no task await/stop).
+
+        Used by reconnection logic to tear down stale state before re-establishing.
+        """
+        self._stdio_error = None
+
+        if self._exit_stack:
+            await self._exit_stack.aclose()
+            self._exit_stack = AsyncExitStack()
+        if self._stdio_task:
+            if self._stdio_stop:
+                self._stdio_stop.set()
+            try:
+                await self._stdio_task
+            except Exception as e:
+                logger.debug("Error stopping stdio task during cleanup: %s", e)
+            self._stdio_task = None
+            self._stdio_ready = None
+            self._stdio_stop = None
+        if self._stdio_exit_stack:
+            await self._stdio_exit_stack.aclose()
+            self._stdio_exit_stack = None
+        self._session = None
+        self._http = None
+        self._write = None
+        self._stdio = None
+        self._get_session_id = None
+        self._session_id = None
+
+    async def _reconnect_session(self) -> None:
+        """Tear down old session and reconnect to MCP server with linear backoff.
+
+        Raises:
+            PluginError: If reconnection fails after all attempts.
+        """
+        logger.info("Attempting to reconnect to MCP server: %s", self.name)
+
+        await self._cleanup_session()
+
+        last_error: Optional[Exception] = None
+        for attempt in range(1, self._reconnect_attempts + 1):
+            try:
+                logger.debug("Reconnection attempt %d/%d to %s", attempt, self._reconnect_attempts, self.name)
+
+                if self._config.mcp.proto == TransportType.STREAMABLEHTTP:
+                    await self.__connect_to_http_server(self._config.mcp.url)
+                elif self._config.mcp.proto == TransportType.STDIO:
+                    await self.__connect_to_stdio_server(
+                        self._config.mcp.script, self._config.mcp.cmd, self._config.mcp.env, self._config.mcp.cwd
+                    )
+
+                logger.info("Reconnected to MCP server on attempt %d: %s", attempt, self.name)
+                return
+            except Exception as e:
+                last_error = e
+                if attempt < self._reconnect_attempts:
+                    delay = self._reconnect_delay * attempt
+                    logger.warning("Reconnection attempt %d failed: %s. Retrying in %ss...", attempt, e, delay)
+                    await asyncio.sleep(delay)
+
+        raise PluginError(
+            error=PluginErrorModel(
+                message=f"Failed to reconnect after {self._reconnect_attempts} attempts: {last_error}",
+                plugin_name=self.name,
+            )
+        )
+
     async def invoke_hook(self, hook_type: str, payload: PluginPayload, context: PluginContext) -> PluginResult:
         """Invoke an external plugin hook using the MCP protocol.
 
@@ -432,11 +506,11 @@ async def invoke_hook(self, hook_type: str, payload: PluginPayload, context: Plu
         if not self._session:
             raise PluginError(error=PluginErrorModel(message="Plugin session not initialized", plugin_name=self.name))
 
-        try:
-            result = await self._session.call_tool(
+        async def _execute_call() -> PluginResult:
+            call_result = await self._session.call_tool(
                 INVOKE_HOOK, {HOOK_TYPE: hook_type, PLUGIN_NAME: self.name, PAYLOAD: payload, CONTEXT: context}
             )
-            for content in result.content:
+            for content in call_result.content:
                 if not isinstance(content, TextContent):
                     continue
                 try:
@@ -457,17 +531,38 @@ async def invoke_hook(self, hook_type: str, payload: PluginPayload, context: Plu
                 if RESULT in res:
                     return result_type.model_validate(res[RESULT])
                 if ERROR in res:
-                    error = PluginErrorModel.model_validate(res[ERROR])
-                    raise PluginError(error)
+                    error_model = PluginErrorModel.model_validate(res[ERROR])
+                    raise PluginError(error_model)
+            raise PluginError(
+                error=PluginErrorModel(
+                    message=f"Received invalid response. Result = {call_result}", plugin_name=self.name
+                )
+            )
+
+        try:
+            return await _execute_call()
         except PluginError as pe:
+            error_msg = str(pe.error.message).lower() if pe.error and pe.error.message else ""
+            if "session" in error_msg and "terminated" in error_msg:
+                logger.warning("Session terminated for plugin %s, attempting reconnection...", self.name)
+                try:
+                    await self._reconnect_session()
+                    return await _execute_call()
+                except Exception as reconn_err:
+                    logger.exception("Reconnection failed for plugin %s: %s", self.name, reconn_err)
             logger.exception(pe)
             raise
+        except McpError as e:
+            logger.warning("McpError for plugin %s: %s", self.name, e)
+            try:
+                await self._reconnect_session()
+                return await _execute_call()
+            except Exception as reconn_err:
+                logger.exception("Reconnection failed for plugin %s: %s", self.name, reconn_err)
+            raise PluginError(error=convert_exception_to_error(e, plugin_name=self.name))
         except Exception as e:
             logger.exception(e)
             raise PluginError(error=convert_exception_to_error(e, plugin_name=self.name))
-        raise PluginError(
-            error=PluginErrorModel(message=f"Received invalid response. Result = {result}", plugin_name=self.name)
-        )
 
     async def __get_plugin_config(self) -> PluginConfig | None:
         """Retrieve plugin configuration for the current plugin on the remote MCP server.

cpex/framework/models.py

@@ -589,6 +589,8 @@ class MCPClientConfig(BaseModel):
         cwd (Optional[str]): Working directory for STDIO server process.
         uds (Optional[str]): Unix domain socket path for streamable HTTP.
         tls (Optional[MCPClientTLSConfig]): Client-side TLS configuration for mTLS.
+        reconnect_attempts (int): Number of reconnection attempts on session failure.
+        reconnect_delay (float): Base delay between reconnection attempts (seconds, linear backoff).
     """
 
     proto: TransportType
@@ -599,6 +601,8 @@ class MCPClientConfig(BaseModel):
     cwd: Optional[str] = None
     uds: Optional[str] = None
     tls: Optional[MCPClientTLSConfig] = None
+    reconnect_attempts: int = Field(default=3, description="Number of reconnection attempts on session failure")
+    reconnect_delay: float = Field(default=0.1, description="Base delay between reconnection attempts (seconds)")
 
     @field_validator(URL, mode="after")
     @classmethod