Skip to content

Python: Make MCP connection-loss recovery ownership-aware #9583

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jose Alvarez (jpalvarezl) merged 3 commits into from
Oct 22, 2026
+302 −84
Merged

Python: Make MCP connection-loss recovery ownership-aware #9583

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
18d099a
Python: Make MCP recovery ownership-aware
jpalvarezl Oct 21, 2026
02146c7
Refactor MCP connection recovery
jpalvarezl Oct 21, 2026
64bf47f
Address MCP recovery review feedback
jpalvarezl Oct 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 111 additions & 72 deletions python/packages/core/agent_framework/_mcp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2279,6 +2279,44 @@ def _set_server_capabilities(self, capabilities: types.ServerCapabilities | None
self._supports_prompts = getattr(capabilities, "prompts", None) is not None
self._supports_logging = getattr(capabilities, "logging", None) is not None

async def _recover_connection_after_loss(
self,
*,
operation: str,
inner_exception: Exception,
load_configured: bool,
) -> None:
"""Recreate a framework-owned connection or fail clearly for a caller-owned connection."""
connection = self._connection
if connection is not None and not connection.is_framework_owned:
connection_type = "Client" if connection.client is not None else "ClientSession"
raise _mcp_tool_execution_exception(
f"MCP connection lost during {operation}. The active {connection_type} is caller-owned and cannot "
"be recreated by Agent Framework. Create a replacement connection and MCP tool instance.",
inner_exception=inner_exception,
) from inner_exception

logger.info("MCP connection was lost during %s. Reconnecting...", operation)
try:
if load_configured:
await self.connect(reset=True)
else:
await self._reconnect_without_loading()
except Exception as reconnect_ex:
raise _mcp_tool_execution_exception(
f"Failed to reconnect to MCP server after {operation}.",
inner_exception=reconnect_ex,
) from reconnect_ex

def _is_connection_loss(self, exception: Exception) -> bool:
"""Return whether an exception proves that the active MCP connection ended."""
from anyio import ClosedResourceError
from mcp import MCPError

return isinstance(exception, ClosedResourceError) or (
isinstance(exception, MCPError) and "session terminated" in exception.error.message.lower()
)

async def _reconnect_without_loading(self) -> None:
if self._is_lifecycle_owner_task():
await self._connect_on_owner(reset=True, load_configured=False)
Expand Down Expand Up @@ -2903,14 +2941,11 @@ async def _load_prompts_locked(self) -> None:
break
except ClosedResourceError as cl_ex:
if attempt == 0:
logger.info("MCP connection closed unexpectedly while loading prompts. Reconnecting...")
try:
await self._reconnect_without_loading()
except Exception as reconn_ex:
raise _mcp_tool_execution_exception(
"Failed to reconnect to MCP server.",
inner_exception=reconn_ex,
) from reconn_ex
await self._recover_connection_after_loss(
operation="prompt discovery",
inner_exception=cl_ex,
load_configured=False,
)
continue
_log_mcp_exception(
logging.ERROR,
Expand Down Expand Up @@ -3049,14 +3084,11 @@ async def _load_tools_locked(self) -> None:
break
except ClosedResourceError as cl_ex:
if attempt == 0:
logger.info("MCP connection closed unexpectedly while loading tools. Reconnecting...")
try:
await self._reconnect_without_loading()
except Exception as reconn_ex:
raise _mcp_tool_execution_exception(
"Failed to reconnect to MCP server.",
inner_exception=reconn_ex,
) from reconn_ex
await self._recover_connection_after_loss(
operation="tool discovery",
inner_exception=cl_ex,
load_configured=False,
)
continue
_log_mcp_exception(
logging.ERROR,
Expand Down Expand Up @@ -3234,6 +3266,7 @@ async def _ensure_connected(self) -> None:
Raises:
ToolExecutionException: If reconnection fails.
"""
from anyio import ClosedResourceError
from mcp import MCPError

if not self._ping_available:
Expand All @@ -3246,23 +3279,36 @@ async def _ensure_connected(self) -> None:
self._ping_available = False
logger.debug("Skipping future MCP pings because the server does not support ping.")
return
logger.info("MCP connection invalid or closed. Reconnecting...")
try:
await self._reconnect_without_loading()
except Exception as ex:
raise _mcp_tool_execution_exception(
"Failed to establish MCP connection.",
inner_exception=ex,
) from ex
except Exception:
logger.info("MCP connection invalid or closed. Reconnecting...")
try:
await self._reconnect_without_loading()
except Exception as ex:
raise _mcp_tool_execution_exception(
"Failed to establish MCP connection.",
inner_exception=ex,
) from ex
if self._is_connection_loss(mcp_exc) or self._connection is None or self._connection.is_framework_owned:
await self._recover_connection_after_loss(
operation="connection liveness check",
inner_exception=mcp_exc,
load_configured=False,
)
return
raise _mcp_tool_execution_exception(
"MCP connection liveness check failed.",
inner_exception=mcp_exc,
) from mcp_exc
except ClosedResourceError as connection_ex:
await self._recover_connection_after_loss(
operation="connection liveness check",
inner_exception=connection_ex,
load_configured=False,
)
except Exception as connection_ex:
connection = self._connection
if connection is None or connection.is_framework_owned:
await self._recover_connection_after_loss(
operation="connection liveness check",
inner_exception=connection_ex,
load_configured=False,
)
return
raise _mcp_tool_execution_exception(
"MCP connection liveness check failed.",
inner_exception=connection_ex,
) from connection_ex

async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]:
"""Call a tool with the given arguments.
Expand Down Expand Up @@ -3365,29 +3411,24 @@ async def _call_tool_with_retries(
set_mcp_span_error(span, type(tool_ex).__name__, str(tool_ex))
raise
except (ClosedResourceError, MCPError) as call_ex:
is_session_terminated = (
isinstance(call_ex, MCPError) and "session terminated" in call_ex.error.message.lower()
)
is_connection_lost = isinstance(call_ex, ClosedResourceError) or is_session_terminated
if not is_connection_lost:
if not self._is_connection_loss(call_ex):
error_message = call_ex.error.message if isinstance(call_ex, MCPError) else str(call_ex)
if span.is_recording():
set_mcp_span_error(span, type(call_ex).__name__, error_message)
raise _mcp_tool_execution_exception(error_message, inner_exception=call_ex) from call_ex

if attempt == 0:
# First attempt failed, try reconnecting.
logger.info("MCP connection closed or terminated unexpectedly. Reconnecting...")
try:
await self.connect(reset=True)
continue
except Exception as reconn_ex:
await self._recover_connection_after_loss(
operation="tool call",
inner_exception=call_ex,
load_configured=True,
)
except ToolExecutionException as recovery_error:
if span.is_recording():
set_mcp_span_error(span, type(reconn_ex).__name__, str(reconn_ex))
raise _mcp_tool_execution_exception(
"Failed to reconnect to MCP server.",
inner_exception=reconn_ex,
) from reconn_ex
set_mcp_span_error(span, type(recovery_error).__name__, str(recovery_error))
raise
continue

# Second attempt also failed, give up.
_log_mcp_exception(
Expand Down Expand Up @@ -3508,31 +3549,29 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str:
return parser(prompt_result)
except ClosedResourceError as cl_ex:
if attempt == 0:
# First attempt failed, try reconnecting
logger.info("MCP connection closed unexpectedly. Reconnecting...")
try:
await self.connect(reset=True)
continue # Retry the operation
except Exception as reconn_ex:
set_mcp_span_error(span, type(reconn_ex).__name__, str(reconn_ex))
raise _mcp_tool_execution_exception(
"Failed to reconnect to MCP server.",
inner_exception=reconn_ex,
) from reconn_ex
else:
# Second attempt also failed, give up
_log_mcp_exception(
logging.ERROR,
cl_ex,
"MCP connection closed unexpectedly after reconnection: %s",
cl_ex,
redacted_message="MCP prompt call connection closed after reconnection",
)
set_mcp_span_error(span, type(cl_ex).__name__, str(cl_ex))
raise _mcp_tool_execution_exception(
f"Failed to call prompt '{prompt_name}' - connection lost.",
inner_exception=cl_ex,
) from cl_ex
await self._recover_connection_after_loss(
operation="prompt call",
inner_exception=cl_ex,
load_configured=True,
)
except ToolExecutionException as recovery_error:
set_mcp_span_error(span, type(recovery_error).__name__, str(recovery_error))
raise
continue
# Second attempt also failed, give up
_log_mcp_exception(
logging.ERROR,
cl_ex,
"MCP connection closed unexpectedly after reconnection: %s",
cl_ex,
redacted_message="MCP prompt call connection closed after reconnection",
)
set_mcp_span_error(span, type(cl_ex).__name__, str(cl_ex))
raise _mcp_tool_execution_exception(
f"Failed to call prompt '{prompt_name}' - connection lost.",
inner_exception=cl_ex,
) from cl_ex
except MCPError as mcp_exc:
error_message = mcp_exc.error.message
set_mcp_span_error(span, type(mcp_exc).__name__, error_message)
Expand Down
Loading
Loading