diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 5e8fc27acb6..a7d32007846 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -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) @@ -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, @@ -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, @@ -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: @@ -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. @@ -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( @@ -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) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 0f666f5a618..3128023fffc 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,172 @@ 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]: + 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 != "session": + 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", ["legacy_client", "modern_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("legacy_client" if connection_type == "client" else 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("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: + 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", ["legacy_client", "modern_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() + + +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"), [ @@ -8525,7 +8703,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,11 +8713,11 @@ 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"))), - pytest.raises(ToolExecutionException, match="Failed to establish MCP connection"), + pytest.raises(ToolExecutionException, match="Failed to reconnect to MCP server"), ): await tool._ensure_connected() @@ -8552,7 +8730,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 +8759,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 = [] 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