2121# Third-Party
2222import httpx
2323import orjson
24- from mcp import ClientSession , StdioServerParameters
24+ from mcp import ClientSession , McpError , StdioServerParameters
2525from mcp .client .stdio import stdio_client
2626from mcp .client .streamable_http import streamablehttp_client
2727from mcp .types import TextContent
@@ -85,6 +85,9 @@ def __init__(self, config: PluginConfig) -> None:
8585 self ._get_session_id : Optional [Callable [[], str | None ]] = None
8686 self ._session_id : Optional [str ] = None
8787 self ._http_client_factory : Optional [Callable [..., httpx .AsyncClient ]] = None
88+ self ._reconnect_attempts : int = 3
89+ self ._reconnect_delay : float = 0.1
90+ self ._reconnect_lock : asyncio .Lock = asyncio .Lock ()
8891
8992 async def initialize (self ) -> None :
9093 """Initialize the plugin's connection to the MCP server.
@@ -99,6 +102,10 @@ async def initialize(self) -> None:
99102 message = "The mcp section must be defined for external plugin" , plugin_name = self .name
100103 )
101104 )
105+
106+ self ._reconnect_attempts = self ._config .mcp .reconnect_attempts
107+ self ._reconnect_delay = self ._config .mcp .reconnect_delay
108+
102109 if self ._config .mcp .proto == TransportType .STDIO :
103110 if not (self ._config .mcp .script or self ._config .mcp .cmd ):
104111 raise PluginError (
@@ -405,6 +412,81 @@ def _tls_httpx_client_factory(
405412 logger .info ("Retrying in %ss..." , delay )
406413 await asyncio .sleep (delay )
407414
415+ async def _cleanup_session (self ) -> None :
416+ """Reset session state without a full shutdown (no task await/stop).
417+
418+ Used by reconnection logic to tear down stale state before re-establishing.
419+ """
420+ self ._stdio_error = None
421+
422+ if self ._exit_stack :
423+ await self ._exit_stack .aclose ()
424+ self ._exit_stack = AsyncExitStack ()
425+ if self ._stdio_task :
426+ if self ._stdio_stop :
427+ self ._stdio_stop .set ()
428+ try :
429+ await asyncio .wait_for (self ._stdio_task , timeout = 5.0 )
430+ except asyncio .TimeoutError :
431+ logger .warning ("Stdio task for plugin %s did not exit within 5s, cancelling" , self .name )
432+ self ._stdio_task .cancel ()
433+ try :
434+ await self ._stdio_task
435+ except (asyncio .CancelledError , Exception ):
436+ pass
437+ except Exception as e :
438+ logger .debug ("Error stopping stdio task during cleanup: %s" , e )
439+ self ._stdio_task = None
440+ self ._stdio_ready = None
441+ self ._stdio_stop = None
442+ if self ._stdio_exit_stack :
443+ await self ._stdio_exit_stack .aclose ()
444+ self ._stdio_exit_stack = None
445+ self ._session = None
446+ self ._http = None
447+ self ._write = None
448+ self ._stdio = None
449+ self ._get_session_id = None
450+ self ._session_id = None
451+
452+ async def _reconnect_session (self ) -> None :
453+ """Tear down old session and reconnect to MCP server with linear backoff.
454+
455+ Raises:
456+ PluginError: If reconnection fails after all attempts.
457+ """
458+ logger .info ("Attempting to reconnect to MCP server: %s" , self .name )
459+
460+ await self ._cleanup_session ()
461+
462+ last_error : Optional [Exception ] = None
463+ for attempt in range (1 , self ._reconnect_attempts + 1 ):
464+ try :
465+ logger .debug ("Reconnection attempt %d/%d to %s" , attempt , self ._reconnect_attempts , self .name )
466+
467+ if self ._config .mcp .proto == TransportType .STREAMABLEHTTP :
468+ await self .__connect_to_http_server (self ._config .mcp .url )
469+ elif self ._config .mcp .proto == TransportType .STDIO :
470+ await self .__connect_to_stdio_server (
471+ self ._config .mcp .script , self ._config .mcp .cmd , self ._config .mcp .env , self ._config .mcp .cwd
472+ )
473+
474+ logger .info ("Reconnected to MCP server on attempt %d: %s" , attempt , self .name )
475+ return
476+ except Exception as e :
477+ last_error = e
478+ if attempt < self ._reconnect_attempts :
479+ delay = self ._reconnect_delay * attempt
480+ logger .warning ("Reconnection attempt %d failed: %s. Retrying in %ss..." , attempt , e , delay )
481+ await asyncio .sleep (delay )
482+
483+ raise PluginError (
484+ error = PluginErrorModel (
485+ message = f"Failed to reconnect after { self ._reconnect_attempts } { last_error } ,
486+ plugin_name = self .name ,
487+ )
488+ )
489+
408490 async def invoke_hook (self , hook_type : str , payload : PluginPayload , context : PluginContext ) -> PluginResult :
409491 """Invoke an external plugin hook using the MCP protocol.
410492
@@ -432,11 +514,12 @@ async def invoke_hook(self, hook_type: str, payload: PluginPayload, context: Plu
432514 if not self ._session :
433515 raise PluginError (error = PluginErrorModel (message = "Plugin session not initialized" , plugin_name = self .name ))
434516
435- try :
436- result = await self ._session .call_tool (
517+ async def _execute_call () -> PluginResult :
518+ """Execute the MCP tool call and parse the result."""
519+ call_result = await self ._session .call_tool (
437520 INVOKE_HOOK , {HOOK_TYPE : hook_type , PLUGIN_NAME : self .name , PAYLOAD : payload , CONTEXT : context }
438521 )
439- for content in result .content :
522+ for content in call_result .content :
440523 if not isinstance (content , TextContent ):
441524 continue
442525 try :
@@ -457,17 +540,55 @@ async def invoke_hook(self, hook_type: str, payload: PluginPayload, context: Plu
457540 if RESULT in res :
458541 return result_type .model_validate (res [RESULT ])
459542 if ERROR in res :
460- error = PluginErrorModel .model_validate (res [ERROR ])
461- raise PluginError (error )
543+ error_model = PluginErrorModel .model_validate (res [ERROR ])
544+ raise PluginError (error_model )
545+ raise PluginError (
546+ error = PluginErrorModel (
547+ message = f"Received invalid response. Result = { call_result } , plugin_name = self .name
548+ )
549+ )
550+
551+ try :
552+ return await _execute_call ()
462553 except PluginError as pe :
554+ error_msg = str (pe .error .message ).lower () if pe .error and pe .error .message else ""
555+ if "session" in error_msg and "terminated" in error_msg :
556+ logger .warning ("Session terminated for plugin %s, attempting reconnection..." , self .name )
557+ try :
558+ async with self ._reconnect_lock :
559+ await self ._reconnect_session ()
560+ return await _execute_call ()
561+ except PluginError :
562+ raise
563+ except Exception as reconn_err :
564+ logger .exception ("Reconnection failed for plugin %s: %s" , self .name , reconn_err )
565+ raise PluginError (
566+ error = PluginErrorModel (
567+ message = f"Reconnection failed for plugin { self .name } { reconn_err } ,
568+ plugin_name = self .name ,
569+ )
570+ ) from reconn_err
463571 logger .exception (pe )
464572 raise
573+ except McpError as e :
574+ logger .warning ("McpError for plugin %s: %s" , self .name , e )
575+ try :
576+ async with self ._reconnect_lock :
577+ await self ._reconnect_session ()
578+ return await _execute_call ()
579+ except PluginError :
580+ raise
581+ except Exception as reconn_err :
582+ logger .exception ("Reconnection failed for plugin %s: %s" , self .name , reconn_err )
583+ raise PluginError (
584+ error = PluginErrorModel (
585+ message = f"Reconnection failed for plugin { self .name } { reconn_err } ,
586+ plugin_name = self .name ,
587+ )
588+ ) from reconn_err
465589 except Exception as e :
466590 logger .exception (e )
467591 raise PluginError (error = convert_exception_to_error (e , plugin_name = self .name ))
468- raise PluginError (
469- error = PluginErrorModel (message = f"Received invalid response. Result = { result } , plugin_name = self .name )
470- )
471592
472593 async def __get_plugin_config (self ) -> PluginConfig | None :
473594 """Retrieve plugin configuration for the current plugin on the remote MCP server.
0 commit comments