diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 0052716eaa..6ffa968383 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -267,6 +267,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID and reconnect or effective header-identity changes replace a framework-owned Client's per-Client cache. Agent Framework does not expose shared-store or cache-mode configuration. - **MCP Tasks extension** - The current Agent Framework MCP client does not implement the `io.modelcontextprotocol/tasks` runtime. Tools that advertise `execution.taskSupport == "required"` are skipped with a warning; optional and task-free tools continue through ordinary `tools/call`. Tasks client support is tracked in microsoft/agent-framework#9215. +- **Plain `tools/call` replay gating** - On `ClosedResourceError` / session-terminated `MCPError` the request may already have executed server-side, so `_call_tool_with_retries` only replays after reconnect when the server advertised the tool with `idempotentHint` or `readOnlyHint` and the client opted in with `trust_server_idempotency_hints` (`_is_call_replayable`, backed by the `_tool_annotations_by_name` snapshot from `tools/list`); any other tool raises `ToolExecutionException` stating the outcome is unknown once the connection is back up, with the connection error as the inner exception, and leaves the retry decision to the caller. Eligibility is captured before the first dispatch and the refreshed catalog must preserve it: a reconnect cannot upgrade an originally unhinted operation into a replay, and a replacement that removes the hint also blocks replay. Caller-owned connections keep failing with their existing 'cannot be recreated' error before any of this. Read-only discovery (`tools/list`) and `get_prompt` keep their reconnect-and-retry behavior. ### File Access Harness (`_harness/_file_access.py`) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 37b4ef56cb..50f78e6260 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -1114,6 +1114,7 @@ def __init__( *, max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, tool_result_content: MCPToolResultContentMode = "structured_first", + trust_server_idempotency_hints: bool = False, ) -> None: """Initialize the MCP Tool base. @@ -1163,6 +1164,7 @@ def __init__( self.load_prompts_flag = load_prompts self.parse_prompt_results = parse_prompt_results self.max_host_payload_size_bytes = max_host_payload_size_bytes + self.trust_server_idempotency_hints = trust_server_idempotency_hints self._exit_stack = AsyncExitStack() self._lifecycle_lock = asyncio.Lock() self._lifecycle_request_lock = asyncio.Lock() @@ -1209,6 +1211,7 @@ def __init__( self._progressive_loader_functions: list[FunctionTool] | None = None self._progressive_loaded_tool_names: set[str] = set() self._tool_call_meta_by_name: dict[str, dict[str, Any]] = {} + self._tool_annotations_by_name: dict[str, types.ToolAnnotations | None] = {} self._tool_param_names_by_name: dict[str, set[str]] = {} self._global_extra_arg_names, self._tool_extra_arg_names = _normalize_additional_tool_argument_names( additional_tool_argument_names @@ -2256,6 +2259,7 @@ def _reset_session_discovery_state(self) -> None: if not isinstance((function.additional_properties or {}).get(_MCP_REMOTE_NAME_KEY), str) ] self._tool_call_meta_by_name.clear() + self._tool_annotations_by_name.clear() self._tool_param_names_by_name.clear() self._progressive_loaded_tool_names.clear() @@ -2490,6 +2494,7 @@ async def _connect_on_owner( functions_before_discovery = self._functions.copy() call_meta_before_discovery = self._tool_call_meta_by_name param_names_before_discovery = self._tool_param_names_by_name + annotations_before_discovery = self._tool_annotations_by_name try: connection = self._require_connection() logger.debug("Connected to MCP server: %s", connection.client or connection.session) @@ -2525,6 +2530,7 @@ async def _connect_on_owner( self._functions[:] = functions_before_discovery self._tool_call_meta_by_name = call_meta_before_discovery self._tool_param_names_by_name = param_names_before_discovery + self._tool_annotations_by_name = annotations_before_discovery raise async def _prepare_for_run(self, kwargs: Mapping[str, Any]) -> None: @@ -3212,6 +3218,7 @@ async def _load_tools_locked(self) -> None: self._function_load_callback(function, tool_annotations_by_name[remote_name]) self._functions[:] = current_functions self._tool_call_meta_by_name = tool_call_meta_by_name + self._tool_annotations_by_name = tool_annotations_by_name self._tool_param_names_by_name = tool_param_names_by_name self._progressive_loaded_tool_names.difference_update(existing_tools.keys() - reused_tool_names) @@ -3389,6 +3396,8 @@ async def _call_tool_with_retries( from anyio import ClosedResourceError from mcp import MCPError + # A catalog refresh cannot retroactively make the first dispatch safe to repeat. + replay_eligible = self._is_call_replayable(tool_name) for attempt in range(2): span_error_recorded = False try: @@ -3433,6 +3442,21 @@ async def _call_tool_with_retries( if span.is_recording(): set_mcp_span_error(span, type(recovery_error).__name__, str(recovery_error)) raise + if not replay_eligible or not self._is_call_replayable(tool_name): + # The request may have reached the server before the connection + # dropped, so replaying it could execute a non-idempotent operation + # twice. Surface the unknown outcome instead of a replay; the + # connection is already back up for whoever decides to retry. + if span.is_recording(): + set_mcp_span_error(span, type(call_ex).__name__, str(call_ex)) + raise _mcp_tool_execution_exception( + f"Failed to call tool '{tool_name}' - connection lost and the remote " + "execution outcome is unknown. The connection was re-established but " + "the call was not replayed, because replay requires both a server " + "idempotent/read-only hint and the client-owned " + "trust_server_idempotency_hints opt-in; retry explicitly if safe.", + inner_exception=call_ex, + ) from call_ex continue # Second attempt also failed, give up. @@ -3458,6 +3482,22 @@ async def _call_tool_with_retries( set_mcp_span_error(span, type(retries_exhausted).__name__, str(retries_exhausted)) raise retries_exhausted + def _is_call_replayable(self, tool_name: str) -> bool: + """Return True if *tool_name* may be replayed after connection loss. + + Replay requires both a client-owned trust decision + (``trust_server_idempotency_hints=True``) and a server annotation carrying + ``idempotentHint`` or ``readOnlyHint``. Annotations are supplied by the + server and are untrusted by default, consistent with how the rest of the + package treats ``ToolAnnotations`` as restriction-only hints: without the + client's explicit opt-in, a misannotated or untrusted server could get a + non-idempotent ``tools/call`` executed twice on a lost response. + """ + if not self.trust_server_idempotency_hints: + return False + annotations = self._tool_annotations_by_name.get(tool_name) + return bool(annotations and (annotations.idempotent_hint or annotations.read_only_hint)) + def _resolved_extra_args(self, tool_name: str) -> set[str]: """Return the user-configured extra argument names allowed for a tool.""" return self._global_extra_arg_names | self._tool_extra_arg_names.get(tool_name, set()) @@ -3687,6 +3727,7 @@ def __init__( additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, tool_result_content: MCPToolResultContentMode = "structured_first", + trust_server_idempotency_hints: bool = False, **kwargs: Any, ) -> None: """Initialize the MCP stdio tool. @@ -3794,6 +3835,11 @@ def __init__( transports. ``None`` disables the limit. tool_result_content: How to choose model-visible text when both ``content`` and ``structuredContent`` are present. See :data:`MCPToolResultContentMode`. + trust_server_idempotency_hints: When ``True``, treat server-declared + ``idempotentHint``/``readOnlyHint`` tool annotations as trustworthy, allowing + a ``tools/call`` whose response was lost to a dropped connection to be retried + once after reconnect. Annotations are server-supplied and untrusted by default, + so the default ``False`` never replays an ambiguous call. kwargs: Additional arguments for framework-owned ``StdioServerParameters`` construction. """ if mcp_client is not None and session is not None: @@ -3847,6 +3893,7 @@ def __init__( task_options=task_options, max_host_payload_size_bytes=max_host_payload_size_bytes, tool_result_content=tool_result_content, + trust_server_idempotency_hints=trust_server_idempotency_hints, ) self.command = command self.args = args or [] @@ -3942,6 +3989,7 @@ def __init__( task_options: Any = None, max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, tool_result_content: MCPToolResultContentMode = "structured_first", + trust_server_idempotency_hints: bool = False, **kwargs: Any, ) -> None: """Initialize the MCP streamable HTTP tool. @@ -4120,6 +4168,11 @@ def __init__( transports. ``None`` disables the limit. tool_result_content: How to choose model-visible text when both ``content`` and ``structuredContent`` are present. See :data:`MCPToolResultContentMode`. + trust_server_idempotency_hints: When ``True``, treat server-declared + ``idempotentHint``/``readOnlyHint`` tool annotations as trustworthy, allowing + a ``tools/call`` whose response was lost to a dropped connection to be retried + once after reconnect. Annotations are server-supplied and untrusted by default, + so the default ``False`` never replays an ambiguous call. kwargs: Additional keyword arguments (accepted for backward compatibility but not used). """ if mcp_client is not None and session is not None: @@ -4173,6 +4226,7 @@ def __init__( task_options=task_options, max_host_payload_size_bytes=max_host_payload_size_bytes, tool_result_content=tool_result_content, + trust_server_idempotency_hints=trust_server_idempotency_hints, ) self.url = url self.terminate_on_close = terminate_on_close @@ -4604,6 +4658,7 @@ def __init__( task_options: Any = None, max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, tool_result_content: MCPToolResultContentMode = "structured_first", + trust_server_idempotency_hints: bool = False, **kwargs: Any, ) -> None: """Initialize the deprecated MCP WebSocket compatibility wrapper. @@ -4696,6 +4751,11 @@ def __init__( transports. ``None`` disables the limit. tool_result_content: How to choose model-visible text when both ``content`` and ``structuredContent`` are present. See :data:`MCPToolResultContentMode`. + trust_server_idempotency_hints: When ``True``, treat server-declared + ``idempotentHint``/``readOnlyHint`` tool annotations as trustworthy, allowing + a ``tools/call`` whose response was lost to a dropped connection to be retried + once after reconnect. Annotations are server-supplied and untrusted by default, + so the default ``False`` never replays an ambiguous call. kwargs: Deprecated compatibility arguments. They are not used. """ super().__init__( @@ -4721,6 +4781,7 @@ def __init__( task_options=task_options, max_host_payload_size_bytes=max_host_payload_size_bytes, tool_result_content=tool_result_content, + trust_server_idempotency_hints=trust_server_idempotency_hints, ) self.url = url self._client_kwargs = kwargs diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 3b814aa05d..ba00f5b3df 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -14,7 +14,7 @@ from textwrap import dedent from types import TracebackType from typing import Any, cast -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from httpx2 import AsyncClient, MockTransport, Request, Response @@ -2890,7 +2890,7 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: async def test_mcp_tool_reconnects_after_session_terminated_error(): - """Session termination errors should reconnect once and retry the tool call.""" + """Trusted idempotent tools reconnect once and retry after session termination.""" class TestServer(MCPTool): def __init__(self, **kwargs: Any) -> None: @@ -2914,8 +2914,9 @@ async def connect(self, *, reset: bool = False) -> None: def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: return None # type: ignore[return-value] # pyrefly: ignore[bad-return] # ty: ignore[invalid-return-type] - server = TestServer(name="test_server") + server = TestServer(name="test_server", trust_server_idempotency_hints=True) await server.connect() + server._tool_annotations_by_name["test_tool"] = types.ToolAnnotations(idempotent_hint=True) result = await server.call_tool("test_tool", param="test_value") @@ -2925,6 +2926,196 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: assert server.sessions[1].call_tool.await_count == 1 +async def test_mcp_tool_call_not_replayed_without_idempotency_hint(): + """Tools without an idempotency hint raise an unknown-outcome error instead of replaying. + + The connection can drop after the server executed the request, so re-issuing the same + tools/call could duplicate a non-idempotent side effect. The default path surfaces the + failure without replaying: the connection is re-established, and the caller decides + whether to retry explicitly. + """ + + class TestServer(MCPTool): + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.connect_count = 0 + self.call_tool_mock = AsyncMock(side_effect=MCPError(-32000, "Session terminated")) + + async def connect(self, *, reset: bool = False) -> None: + self.connect_count += 1 + session = Mock(spec=ClientSession) + session.call_tool = self.call_tool_mock + _set_framework_owned_session(self, session) + self.is_connected = True + + def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: + return None # type: ignore[return-value] # pyrefly: ignore[bad-return] # ty: ignore[invalid-return-type] + + server = TestServer(name="test_server") + await server.connect() + + with pytest.raises(ToolExecutionException, match="outcome is unknown") as exc_info: + await server.call_tool("test_tool", param="test_value") + + assert isinstance(exc_info.value.__cause__, MCPError) + assert server.connect_count == 2 + assert server.call_tool_mock.await_count == 1 + + +async def test_mcp_tool_call_not_replayed_when_reconnect_adds_hint(): + """A replacement catalog cannot make an already dispatched operation replayable.""" + + class TestServer(MCPTool): + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.connect_count = 0 + self.sessions: list[Any] = [] + + async def connect(self, *, reset: bool = False) -> None: + self.connect_count += 1 + mock_session = Mock(spec=ClientSession) + _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")) + else: + self._tool_annotations_by_name["test_tool"] = types.ToolAnnotations(idempotent_hint=True) + mock_session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="recovered")]) + ) + self.is_connected = True + + def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: + return None # type: ignore[return-value] # pyrefly: ignore[bad-return] # ty: ignore[invalid-return-type] + + server = TestServer(name="test_server", trust_server_idempotency_hints=True) + await server.connect() + + with pytest.raises(ToolExecutionException, match="outcome is unknown"): + await server.call_tool("test_tool", param="test_value") + + assert server.connect_count == 2 + assert server.sessions[0].call_tool.await_count == 1 + assert server.sessions[1].call_tool.await_count == 0 + + +async def test_mcp_tool_call_not_replayed_on_hint_alone_without_client_trust(): + """A server idempotency hint by itself must not authorize replay. + + ToolAnnotations are server-supplied and untrusted unless the client separately + trusts the server, so with the default trust posture a hinted tool still gets + the unknown-outcome error instead of a replay. + """ + + class TestServer(MCPTool): + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.connect_count = 0 + self.call_tool_mock = AsyncMock(side_effect=MCPError(-32000, "Session terminated")) + + async def connect(self, *, reset: bool = False) -> None: + self.connect_count += 1 + session = Mock(spec=ClientSession) + session.call_tool = self.call_tool_mock + _set_framework_owned_session(self, session) + self.is_connected = True + + def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: + return None # type: ignore[return-value] # pyrefly: ignore[bad-return] # ty: ignore[invalid-return-type] + + server = TestServer(name="test_server") + await server.connect() + server._tool_annotations_by_name["test_tool"] = types.ToolAnnotations(idempotent_hint=True) + + with pytest.raises(ToolExecutionException, match="outcome is unknown"): + await server.call_tool("test_tool", param="test_value") + + assert server.connect_count == 2 + assert server.call_tool_mock.await_count == 1 + + +async def test_mcp_regular_tool_call_does_not_duplicate_side_effect_after_disconnect(): + """A non-idempotent tool whose response is lost mid-call executes at most once.""" + from anyio.streams.memory import ClosedResourceError + + class TestServer(MCPTool): + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.state = 0 + + async def connect(self, *, reset: bool = False) -> None: + session = Mock(spec=ClientSession) + + async def call_tool(*args: Any, **kwargs: Any) -> types.CallToolResult: + self.state += 1 + if self.state == 1: + raise ClosedResourceError + return types.CallToolResult(content=[types.TextContent(type="text", text="ok")]) + + session.call_tool = call_tool + _set_framework_owned_session(self, session) + self.is_connected = True + + def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: + return None # type: ignore[return-value] # pyrefly: ignore[bad-return] # ty: ignore[invalid-return-type] + + server = TestServer(name="test_server") + await server.connect() + + with pytest.raises(ToolExecutionException, match="outcome is unknown"): + await server.call_tool("test_tool", param="test_value") + + assert server.state == 1 + + +async def test_connect_discovery_failure_rolls_back_annotations_snapshot(): + """When load_tools succeeds but a later discovery step fails, the rollback in + _connect_on_owner must restore the pre-discovery annotations map along with the + other metadata, or a later replay decision reads hints from the failed session.""" + tool = MCPStdioTool(name="test", command="test-command") + sentinel_annotations = {"prev_tool": types.ToolAnnotations(read_only_hint=True)} + sentinel_meta = {"prev_tool": {"version": 0}} + sentinel_params = {"prev_tool": {"q"}} + tool._tool_annotations_by_name = dict(sentinel_annotations) + tool._tool_call_meta_by_name = dict(sentinel_meta) + tool._tool_param_names_by_name = dict(sentinel_params) + retained = FunctionTool(name="retained", func=lambda: "ok") + tool._functions.append(retained) + + mock_transport = (MagicMock(), MagicMock()) + mock_context_manager = Mock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport) + mock_context_manager.__aexit__ = AsyncMock(return_value=None) + tool.get_mcp_client = Mock(return_value=mock_context_manager) # type: ignore[method-assign] + + sdk_client = _mock_sdk_client( + capabilities=types.ServerCapabilities( + tools=types.ToolsCapability(list_changed=False), + prompts=types.PromptsCapability(list_changed=False), + ) + ) + sdk_client.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="fresh_tool", + input_schema={"type": "object", "properties": {}}, + annotations=types.ToolAnnotations(idempotent_hint=True), + ) + ] + ) + ) + sdk_client.list_prompts = AsyncMock(side_effect=RuntimeError("prompt boom")) + + with patch("mcp.Client", return_value=sdk_client), pytest.raises(RuntimeError, match="prompt boom"): + await tool.connect() + + assert tool._tool_annotations_by_name == sentinel_annotations + assert tool._tool_call_meta_by_name == sentinel_meta + assert tool._tool_param_names_by_name == sentinel_params + assert tool._functions == [retained] + + async def test_mcp_tool_call_tool_raises_on_is_error(): """Test that call_tool raises ToolExecutionException when MCP returns is_error=True.""" @@ -3978,7 +4169,15 @@ async def test_mcp_connection_reset_integration(): """ url = os.environ.get("LOCAL_MCP_URL") - tool = MCPStreamableHTTPTool(name="integration_test", url=url, approval_mode="never_require") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + # The query tool below is annotated read-only before the simulated drop; + # replay also needs the client-owned trust opt-in, which the default + # posture deliberately withholds. + tool = MCPStreamableHTTPTool( + name="integration_test", + url=url, + approval_mode="never_require", # type: ignore[arg-type] + trust_server_idempotency_hints=True, + ) # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] async with tool: # Verify initial connection @@ -3997,6 +4196,10 @@ async def test_mcp_connection_reset_integration(): original_exit_stack = tool._exit_stack original_call_tool = tool.session.call_tool + # The query tool is read-only, so a lost response is safe to replay: mark it as such. + remote_name = (func.additional_properties or {}).get("_mcp_remote_name") or func.name + tool._tool_annotations_by_name[remote_name] = types.ToolAnnotations(read_only_hint=True) + # Simulate connection failure by making call_tool raise ClosedResourceError once call_count = 0 @@ -7044,6 +7247,9 @@ async def test_mcp_tool_connection_properly_invalidated_after_closed_resource_er This test verifies the fix for issue #2884: the tool tries operations optimistically and only reconnects when ClosedResourceError is encountered, avoiding extra latency. + The replay only happens because the client opted into trusting the tool's idempotency + hint here; tools without that hint, or without the client opt-in, raise instead of + retrying (issue #9204). """ from unittest.mock import AsyncMock, MagicMock, patch @@ -7058,7 +7264,9 @@ async def test_mcp_tool_connection_properly_invalidated_after_closed_resource_er command="test_command", args=["arg1"], load_tools=True, + trust_server_idempotency_hints=True, ) + tool._tool_annotations_by_name["test_tool"] = types.ToolAnnotations(idempotent_hint=True) # Mock the session mock_session = MagicMock() @@ -7299,11 +7507,14 @@ async def test_mcp_tool_get_prompt_requires_loaded_prompts() -> None: async def test_mcp_tool_call_tool_raises_after_reconnection_still_fails() -> None: from anyio.streams.memory import ClosedResourceError - tool = MCPTool(name="test_tool", load_tools=True) # type: ignore[abstract] # ty: ignore[call-non-callable] + tool = MCPTool( # type: ignore[abstract] # ty: ignore[call-non-callable] + name="test_tool", load_tools=True, trust_server_idempotency_hints=True + ) _set_framework_owned_session( tool, Mock(call_tool=AsyncMock(side_effect=[ClosedResourceError(), ClosedResourceError()])), ) + tool._tool_annotations_by_name["remote_tool"] = types.ToolAnnotations(idempotent_hint=True) with ( patch.object(tool, "connect", AsyncMock()) as mock_connect, diff --git a/python/packages/core/tests/core/test_mcp_observability.py b/python/packages/core/tests/core/test_mcp_observability.py index 9858823119..4115c5abd4 100644 --- a/python/packages/core/tests/core/test_mcp_observability.py +++ b/python/packages/core/tests/core/test_mcp_observability.py @@ -327,6 +327,8 @@ async def test_mcp_tools_call_second_connection_loss_log_respects_sensitive_data """Second connection loss should retain an error log without exporting exception details by default.""" private_error_detail = "founder-secret-token" tool = _make_connected_mcp_tool() + tool.trust_server_idempotency_hints = True + tool._tool_annotations_by_name["get-weather"] = types.ToolAnnotations(idempotent_hint=True) tool.session.call_tool = AsyncMock( # type: ignore[method-assign, union-attr] # ty: ignore[invalid-assignment] side_effect=[ MCPError(-32000, f"session terminated: {private_error_detail}"), @@ -342,6 +344,8 @@ async def test_mcp_tools_call_second_connection_loss_log_respects_sensitive_data ): await tool.call_tool("get-weather") + assert tool.session is not None + assert getattr(tool.session.call_tool, "await_count", 0) == 2 assert (private_error_detail in caplog.text) is enable_sensitive_data if not enable_sensitive_data: assert "MCP tool call connection closed after reconnection (MCPError)." in caplog.text