From 18d099a5113db4936377e54f75e1b5e682e251a5 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Wed, 21 Oct 2026 18:18:22 +0200 Subject: [PATCH 1/3] Python: Make MCP recovery ownership-aware --- python/packages/core/agent_framework/_mcp.py | 59 +++++++- python/packages/core/tests/core/test_mcp.py | 151 +++++++++++++++++-- 2 files changed, 198 insertions(+), 12 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 5e8fc27acb6..96eb36ef16c 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -2279,6 +2279,24 @@ 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 + def _caller_owned_connection_loss_error( + self, + *, + operation: str, + inner_exception: Exception, + ) -> ToolExecutionException | None: + """Return an actionable error when Agent Framework cannot recreate the active connection.""" + connection = self._require_connection() + if connection.is_framework_owned: + return None + + connection_type = "Client" if connection.client is not None else "ClientSession" + return _mcp_tool_execution_exception( + f"MCP connection was 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, + ) + async def _reconnect_without_loading(self) -> None: if self._is_lifecycle_owner_task(): await self._connect_on_owner(reset=True, load_configured=False) @@ -2903,6 +2921,12 @@ async def _load_prompts_locked(self) -> None: break except ClosedResourceError as cl_ex: if attempt == 0: + caller_owned_error = self._caller_owned_connection_loss_error( + operation="prompt discovery", + inner_exception=cl_ex, + ) + if caller_owned_error is not None: + raise caller_owned_error from cl_ex logger.info("MCP connection closed unexpectedly while loading prompts. Reconnecting...") try: await self._reconnect_without_loading() @@ -3049,6 +3073,12 @@ async def _load_tools_locked(self) -> None: break except ClosedResourceError as cl_ex: if attempt == 0: + caller_owned_error = self._caller_owned_connection_loss_error( + operation="tool discovery", + inner_exception=cl_ex, + ) + if caller_owned_error is not None: + raise caller_owned_error from cl_ex logger.info("MCP connection closed unexpectedly while loading tools. Reconnecting...") try: await self._reconnect_without_loading() @@ -3246,6 +3276,12 @@ async def _ensure_connected(self) -> None: self._ping_available = False logger.debug("Skipping future MCP pings because the server does not support ping.") return + caller_owned_error = self._caller_owned_connection_loss_error( + operation="connection liveness check", + inner_exception=mcp_exc, + ) + if caller_owned_error is not None: + raise caller_owned_error from mcp_exc logger.info("MCP connection invalid or closed. Reconnecting...") try: await self._reconnect_without_loading() @@ -3254,7 +3290,13 @@ async def _ensure_connected(self) -> None: "Failed to establish MCP connection.", inner_exception=ex, ) from ex - except Exception: + except Exception as connection_ex: + caller_owned_error = self._caller_owned_connection_loss_error( + operation="connection liveness check", + inner_exception=connection_ex, + ) + if caller_owned_error is not None: + raise caller_owned_error from connection_ex logger.info("MCP connection invalid or closed. Reconnecting...") try: await self._reconnect_without_loading() @@ -3376,6 +3418,14 @@ async def _call_tool_with_retries( raise _mcp_tool_execution_exception(error_message, inner_exception=call_ex) from call_ex if attempt == 0: + caller_owned_error = self._caller_owned_connection_loss_error( + operation=f"tool call '{tool_name}'", + inner_exception=call_ex, + ) + if caller_owned_error is not None: + if span.is_recording(): + set_mcp_span_error(span, type(caller_owned_error).__name__, str(caller_owned_error)) + raise caller_owned_error from call_ex # First attempt failed, try reconnecting. logger.info("MCP connection closed or terminated unexpectedly. Reconnecting...") try: @@ -3508,6 +3558,13 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: return parser(prompt_result) except ClosedResourceError as cl_ex: if attempt == 0: + caller_owned_error = self._caller_owned_connection_loss_error( + operation=f"prompt call '{prompt_name}'", + inner_exception=cl_ex, + ) + if caller_owned_error is not None: + set_mcp_span_error(span, type(caller_owned_error).__name__, str(caller_owned_error)) + raise caller_owned_error from cl_ex # First attempt failed, try reconnecting logger.info("MCP connection closed unexpectedly. Reconnecting...") try: diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 0f666f5a618..1078bf06abc 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -164,6 +164,12 @@ def _connected_client(tool: MCPTool) -> Any: return connection.client +def _set_framework_owned_session(tool: MCPTool, session: ClientSession) -> None: + """Model the pre-v2 framework-owned session fixtures used by reconnect tests.""" + tool.session = session + tool._require_connection().is_framework_owned = True + + class _TransportBoundClientContext: """Model SDK-owned transport entry and exit without running a protocol dispatcher.""" @@ -2885,7 +2891,7 @@ def __init__(self, **kwargs: Any) -> None: async def connect(self, *, reset: bool = False) -> None: self.connect_count += 1 mock_session = Mock(spec=ClientSession) - self.session = mock_session + _set_framework_owned_session(self, mock_session) self.sessions.append(self.session) if self.connect_count == 1: mock_session.call_tool = AsyncMock(side_effect=MCPError(-32000, "Session terminated")) @@ -7041,14 +7047,14 @@ async def test_mcp_tool_connection_properly_invalidated_after_closed_resource_er async def restore_session(*, reset=False): if reset: await original_exit_stack.aclose() - tool.session = mock_session + _set_framework_owned_session(tool, mock_session) tool.is_connected = True tool._tools_loaded = True mock_connect.side_effect = restore_session # Simulate initial connection - tool.session = mock_session + _set_framework_owned_session(tool, mock_session) tool.is_connected = True tool._tools_loaded = True @@ -7139,14 +7145,14 @@ async def test_mcp_tool_get_prompt_reconnection_on_closed_resource_error(): async def restore_session(*, reset=False): if reset: await original_exit_stack.aclose() - tool.session = mock_session + _set_framework_owned_session(tool, mock_session) tool.is_connected = True tool._prompts_loaded = True mock_connect.side_effect = restore_session # Simulate initial connection - tool.session = mock_session + _set_framework_owned_session(tool, mock_session) tool.is_connected = True tool._prompts_loaded = True @@ -7267,7 +7273,10 @@ async def test_mcp_tool_call_tool_raises_after_reconnection_still_fails() -> Non from anyio.streams.memory import ClosedResourceError tool = MCPTool(name="test_tool", load_tools=True) # type: ignore[abstract] # ty: ignore[call-non-callable] - tool.session = Mock(call_tool=AsyncMock(side_effect=[ClosedResourceError(), ClosedResourceError()])) + _set_framework_owned_session( + tool, + Mock(call_tool=AsyncMock(side_effect=[ClosedResourceError(), ClosedResourceError()])), + ) with ( patch.object(tool, "connect", AsyncMock()) as mock_connect, @@ -7284,7 +7293,10 @@ async def test_mcp_tool_get_prompt_raises_after_reconnection_still_fails() -> No from anyio.streams.memory import ClosedResourceError tool = MCPTool(name="test_tool", load_prompts=True) # type: ignore[abstract] # ty: ignore[call-non-callable] - tool.session = Mock(get_prompt=AsyncMock(side_effect=[ClosedResourceError(), ClosedResourceError()])) + _set_framework_owned_session( + tool, + Mock(get_prompt=AsyncMock(side_effect=[ClosedResourceError(), ClosedResourceError()])), + ) with ( patch.object(tool, "connect", AsyncMock()) as mock_connect, @@ -8298,6 +8310,123 @@ async def test_mcp_tool_reset_preserves_supplied_client() -> None: assert client.session is session +def _caller_owned_mcp_tool(connection_type: str) -> tuple[MCPStdioTool, Any]: + client = _mock_sdk_client() + session = client.session + + if connection_type == "client": + return ( + MCPStdioTool( + name="test", + command="unused", + mcp_client=client, + load_tools=True, + load_prompts=True, + ), + client, + ) + + return ( + MCPStdioTool( + name="test", + command="unused", + session=session, + load_tools=True, + load_prompts=True, + ), + session, + ) + + +@pytest.mark.parametrize("connection_type", ["client", "session"]) +@pytest.mark.parametrize( + ("operation_name", "load_method", "capability_attribute"), + [ + ("list_tools", "load_tools", "_supports_tools"), + ("list_prompts", "load_prompts", "_supports_prompts"), + ], +) +async def test_caller_owned_connection_loss_during_discovery_does_not_reconnect( + connection_type: str, + operation_name: str, + load_method: str, + capability_attribute: str, +) -> None: + from anyio import ClosedResourceError + + tool, connection = _caller_owned_mcp_tool(connection_type) + operation = AsyncMock(side_effect=ClosedResourceError()) + setattr(connection, operation_name, operation) + setattr(tool, capability_attribute, True) + + with ( + patch.object(tool, "_reconnect_without_loading", new_callable=AsyncMock) as reconnect, + pytest.raises(ToolExecutionException, match=r"caller-owned and cannot be recreated"), + ): + await getattr(tool, load_method)() + + operation.assert_awaited_once() + reconnect.assert_not_awaited() + + +@pytest.mark.parametrize("connection_type", ["client", "session"]) +async def test_caller_owned_connection_loss_during_liveness_check_does_not_reconnect( + connection_type: str, +) -> None: + from anyio import ClosedResourceError + + tool, connection = _caller_owned_mcp_tool(connection_type) + connection.send_ping = AsyncMock(side_effect=ClosedResourceError()) + tool._ping_available = True + + with ( + patch.object(tool, "_reconnect_without_loading", new_callable=AsyncMock) as reconnect, + pytest.raises(ToolExecutionException, match=r"caller-owned and cannot be recreated"), + ): + await tool._ensure_connected() + + connection.send_ping.assert_awaited_once() + reconnect.assert_not_awaited() + + +@pytest.mark.parametrize("connection_type", ["client", "session"]) +async def test_caller_owned_connection_loss_during_tool_call_does_not_reconnect( + connection_type: str, +) -> None: + from anyio import ClosedResourceError + + tool, connection = _caller_owned_mcp_tool(connection_type) + connection.call_tool = AsyncMock(side_effect=ClosedResourceError()) + + with ( + patch.object(tool, "connect", new_callable=AsyncMock) as reconnect, + pytest.raises(ToolExecutionException, match=r"caller-owned and cannot be recreated"), + ): + await tool.call_tool("mutating_tool") + + connection.call_tool.assert_awaited_once() + reconnect.assert_not_awaited() + + +@pytest.mark.parametrize("connection_type", ["client", "session"]) +async def test_caller_owned_connection_loss_during_prompt_call_does_not_reconnect( + connection_type: str, +) -> None: + from anyio import ClosedResourceError + + tool, connection = _caller_owned_mcp_tool(connection_type) + connection.get_prompt = AsyncMock(side_effect=ClosedResourceError()) + + with ( + patch.object(tool, "connect", new_callable=AsyncMock) as reconnect, + pytest.raises(ToolExecutionException, match=r"caller-owned and cannot be recreated"), + ): + await tool.get_prompt("test_prompt") + + connection.get_prompt.assert_awaited_once() + reconnect.assert_not_awaited() + + @pytest.mark.parametrize( ("mode", "expected_version"), [ @@ -8525,7 +8654,7 @@ async def test_ensure_connected_skips_future_pings_when_ping_is_not_available() async def test_ensure_connected_reconnects_on_failed_ping() -> None: tool = MCPTool(name="test_tool") # type: ignore[abstract] # ty: ignore[call-non-callable] - tool.session = Mock(send_ping=AsyncMock(side_effect=RuntimeError("closed"))) + _set_framework_owned_session(tool, Mock(send_ping=AsyncMock(side_effect=RuntimeError("closed")))) with patch.object(tool, "_reconnect_without_loading", AsyncMock()) as mock_reconnect: await tool._ensure_connected() @@ -8535,7 +8664,7 @@ async def test_ensure_connected_reconnects_on_failed_ping() -> None: async def test_ensure_connected_wraps_reconnect_failure() -> None: tool = MCPTool(name="test_tool") # type: ignore[abstract] # ty: ignore[call-non-callable] - tool.session = Mock(send_ping=AsyncMock(side_effect=RuntimeError("closed"))) + _set_framework_owned_session(tool, Mock(send_ping=AsyncMock(side_effect=RuntimeError("closed")))) with ( patch.object(tool, "_reconnect_without_loading", AsyncMock(side_effect=RuntimeError("still closed"))), @@ -8552,7 +8681,7 @@ async def test_load_tools_reconnects_on_closed_resource_when_ping_is_unavailable first_session = Mock() first_session.list_tools = AsyncMock(side_effect=ClosedResourceError()) - tool.session = first_session + _set_framework_owned_session(tool, first_session) page = Mock() page.tools = [] @@ -8581,7 +8710,7 @@ async def test_load_prompts_reconnects_on_closed_resource_when_ping_is_unavailab first_session = Mock() first_session.list_prompts = AsyncMock(side_effect=ClosedResourceError()) - tool.session = first_session + _set_framework_owned_session(tool, first_session) page = Mock() page.prompts = [] From 02146c7d840fbeb425721833026746c78b66cc87 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Wed, 21 Oct 2026 18:23:52 +0200 Subject: [PATCH 2/3] Refactor MCP connection recovery --- python/packages/core/agent_framework/_mcp.py | 165 +++++++------------ python/packages/core/tests/core/test_mcp.py | 2 +- 2 files changed, 62 insertions(+), 105 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 96eb36ef16c..b784b70f31a 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -2279,23 +2279,34 @@ 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 - def _caller_owned_connection_loss_error( + async def _recover_connection_after_loss( self, *, operation: str, inner_exception: Exception, - ) -> ToolExecutionException | None: - """Return an actionable error when Agent Framework cannot recreate the active connection.""" + load_configured: bool, + ) -> None: + """Recreate a framework-owned connection or fail clearly for a caller-owned connection.""" connection = self._require_connection() - if connection.is_framework_owned: - return None + if not connection.is_framework_owned: + connection_type = "Client" if connection.client is not None else "ClientSession" + raise _mcp_tool_execution_exception( + f"MCP connection was 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 - connection_type = "Client" if connection.client is not None else "ClientSession" - return _mcp_tool_execution_exception( - f"MCP connection was 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, - ) + 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 async def _reconnect_without_loading(self) -> None: if self._is_lifecycle_owner_task(): @@ -2921,20 +2932,11 @@ async def _load_prompts_locked(self) -> None: break except ClosedResourceError as cl_ex: if attempt == 0: - caller_owned_error = self._caller_owned_connection_loss_error( + await self._recover_connection_after_loss( operation="prompt discovery", inner_exception=cl_ex, + load_configured=False, ) - if caller_owned_error is not None: - raise caller_owned_error from cl_ex - 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 continue _log_mcp_exception( logging.ERROR, @@ -3073,20 +3075,11 @@ async def _load_tools_locked(self) -> None: break except ClosedResourceError as cl_ex: if attempt == 0: - caller_owned_error = self._caller_owned_connection_loss_error( + await self._recover_connection_after_loss( operation="tool discovery", inner_exception=cl_ex, + load_configured=False, ) - if caller_owned_error is not None: - raise caller_owned_error from cl_ex - 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 continue _log_mcp_exception( logging.ERROR, @@ -3276,35 +3269,17 @@ async def _ensure_connected(self) -> None: self._ping_available = False logger.debug("Skipping future MCP pings because the server does not support ping.") return - caller_owned_error = self._caller_owned_connection_loss_error( + await self._recover_connection_after_loss( operation="connection liveness check", inner_exception=mcp_exc, + load_configured=False, ) - if caller_owned_error is not None: - raise caller_owned_error from mcp_exc - 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 as connection_ex: - caller_owned_error = self._caller_owned_connection_loss_error( + await self._recover_connection_after_loss( operation="connection liveness check", inner_exception=connection_ex, + load_configured=False, ) - if caller_owned_error is not None: - raise caller_owned_error from connection_ex - 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 async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: """Call a tool with the given arguments. @@ -3418,26 +3393,17 @@ async def _call_tool_with_retries( raise _mcp_tool_execution_exception(error_message, inner_exception=call_ex) from call_ex if attempt == 0: - caller_owned_error = self._caller_owned_connection_loss_error( - operation=f"tool call '{tool_name}'", - inner_exception=call_ex, - ) - if caller_owned_error is not None: - if span.is_recording(): - set_mcp_span_error(span, type(caller_owned_error).__name__, str(caller_owned_error)) - raise caller_owned_error from call_ex - # 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( @@ -3558,38 +3524,29 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: return parser(prompt_result) except ClosedResourceError as cl_ex: if attempt == 0: - caller_owned_error = self._caller_owned_connection_loss_error( - operation=f"prompt call '{prompt_name}'", - inner_exception=cl_ex, - ) - if caller_owned_error is not None: - set_mcp_span_error(span, type(caller_owned_error).__name__, str(caller_owned_error)) - raise caller_owned_error from cl_ex - # 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) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 1078bf06abc..f99298a1e7e 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -8668,7 +8668,7 @@ async def test_ensure_connected_wraps_reconnect_failure() -> None: with ( patch.object(tool, "_reconnect_without_loading", AsyncMock(side_effect=RuntimeError("still closed"))), - pytest.raises(ToolExecutionException, match="Failed to establish MCP connection"), + pytest.raises(ToolExecutionException, match="Failed to reconnect to MCP server"), ): await tool._ensure_connected() From 64bf47f992cb3c98ed9769cc0fb99d5a4093b5e2 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Wed, 21 Oct 2026 18:56:30 +0200 Subject: [PATCH 3/3] Address MCP recovery review feedback --- python/packages/core/agent_framework/_mcp.py | 51 ++++++++++++---- python/packages/core/tests/core/test_mcp.py | 61 +++++++++++++++++-- .../core/tests/core/test_mcp_observability.py | 1 + 3 files changed, 94 insertions(+), 19 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index b784b70f31a..a7d32007846 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -2287,11 +2287,11 @@ async def _recover_connection_after_loss( load_configured: bool, ) -> None: """Recreate a framework-owned connection or fail clearly for a caller-owned connection.""" - connection = self._require_connection() - if not connection.is_framework_owned: + 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 was lost during {operation}. The active {connection_type} is caller-owned and cannot " + 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 @@ -2308,6 +2308,15 @@ async def _recover_connection_after_loss( 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) @@ -3257,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: @@ -3269,17 +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 - await self._recover_connection_after_loss( - operation="connection liveness check", + 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, - load_configured=False, - ) - except Exception as connection_ex: + ) 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. @@ -3382,11 +3411,7 @@ 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) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index f99298a1e7e..3128023fffc 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -8311,10 +8311,11 @@ async def test_mcp_tool_reset_preserves_supplied_client() -> None: def _caller_owned_mcp_tool(connection_type: str) -> tuple[MCPStdioTool, Any]: - client = _mock_sdk_client() + protocol_version = "2026-07-28" if connection_type == "modern_client" else "2025-11-25" + client = _mock_sdk_client(protocol_version=protocol_version) session = client.session - if connection_type == "client": + if connection_type != "session": return ( MCPStdioTool( name="test", @@ -8338,7 +8339,7 @@ def _caller_owned_mcp_tool(connection_type: str) -> tuple[MCPStdioTool, Any]: ) -@pytest.mark.parametrize("connection_type", ["client", "session"]) +@pytest.mark.parametrize("connection_type", ["legacy_client", "modern_client", "session"]) @pytest.mark.parametrize( ("operation_name", "load_method", "capability_attribute"), [ @@ -8375,7 +8376,7 @@ async def test_caller_owned_connection_loss_during_liveness_check_does_not_recon ) -> None: from anyio import ClosedResourceError - tool, connection = _caller_owned_mcp_tool(connection_type) + tool, connection = _caller_owned_mcp_tool("legacy_client" if connection_type == "client" else connection_type) connection.send_ping = AsyncMock(side_effect=ClosedResourceError()) tool._ping_available = True @@ -8389,7 +8390,34 @@ async def test_caller_owned_connection_loss_during_liveness_check_does_not_recon reconnect.assert_not_awaited() -@pytest.mark.parametrize("connection_type", ["client", "session"]) +@pytest.mark.parametrize("failure", [MCPError(-32000, "server busy"), RuntimeError("temporary failure")]) +async def test_ambiguous_caller_owned_liveness_failure_does_not_require_replacement( + failure: Exception, +) -> None: + tool, connection = _caller_owned_mcp_tool("legacy_client") + connection.send_ping = AsyncMock(side_effect=failure) + tool._ping_available = True + + with ( + patch.object(tool, "_reconnect_without_loading", new_callable=AsyncMock) as reconnect, + pytest.raises(ToolExecutionException, match=r"liveness check failed"), + ): + await tool._ensure_connected() + + connection.send_ping.assert_awaited_once() + reconnect.assert_not_awaited() + + +async def test_modern_caller_owned_liveness_check_does_not_ping() -> None: + tool, connection = _caller_owned_mcp_tool("modern_client") + connection.send_ping = AsyncMock() + + await tool._ensure_connected() + + connection.send_ping.assert_not_awaited() + + +@pytest.mark.parametrize("connection_type", ["legacy_client", "modern_client", "session"]) async def test_caller_owned_connection_loss_during_tool_call_does_not_reconnect( connection_type: str, ) -> None: @@ -8408,7 +8436,7 @@ async def test_caller_owned_connection_loss_during_tool_call_does_not_reconnect( reconnect.assert_not_awaited() -@pytest.mark.parametrize("connection_type", ["client", "session"]) +@pytest.mark.parametrize("connection_type", ["legacy_client", "modern_client", "session"]) async def test_caller_owned_connection_loss_during_prompt_call_does_not_reconnect( connection_type: str, ) -> None: @@ -8427,6 +8455,27 @@ async def test_caller_owned_connection_loss_during_prompt_call_does_not_reconnec reconnect.assert_not_awaited() +async def test_framework_owned_discovery_reconnects_before_connect_and_after_close() -> None: + tool = MCPTool(name="test", load_tools=True, load_prompts=True) # type: ignore[abstract] # ty: ignore[call-non-callable] + tools_session = Mock(spec=ClientSession) + tools_session.list_tools = AsyncMock(return_value=types.ListToolsResult(tools=[])) + prompts_session = Mock(spec=ClientSession) + prompts_session.list_prompts = AsyncMock(return_value=types.ListPromptsResult(prompts=[])) + replacement_sessions = iter([tools_session, prompts_session]) + + async def reconnect() -> None: + _set_framework_owned_session(tool, next(replacement_sessions)) + + with patch.object(tool, "_reconnect_without_loading", AsyncMock(side_effect=reconnect)) as mock_reconnect: + await tool.load_tools() + await tool.close() + await tool.load_prompts() + + assert mock_reconnect.await_count == 2 + tools_session.list_tools.assert_awaited_once() + prompts_session.list_prompts.assert_awaited_once() + + @pytest.mark.parametrize( ("mode", "expected_version"), [ diff --git a/python/packages/core/tests/core/test_mcp_observability.py b/python/packages/core/tests/core/test_mcp_observability.py index 9af3d1c7655..17dd3f62ef9 100644 --- a/python/packages/core/tests/core/test_mcp_observability.py +++ b/python/packages/core/tests/core/test_mcp_observability.py @@ -33,6 +33,7 @@ def _make_connected_mcp_tool( """Create an MCPTool with a mocked session, ready for testing.""" tool = MCPTool(name=name) # type: ignore[abstract] # ty: ignore[call-non-callable] tool.session = AsyncMock() + tool._require_connection().is_framework_owned = True tool.is_connected = True tool._supports_tools = supports_tools tool._supports_prompts = supports_prompts