diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 895764eb93..271639d32e 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -242,6 +242,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** - MCP Python SDK 2.2 does not implement the 2026 `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`. Restoration of the experimental Tasks API and removed sample is tracked in microsoft/agent-framework#8245. +- **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. 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 c4f93237f0..e227625029 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. @@ -1161,6 +1162,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() @@ -1207,6 +1209,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 @@ -2254,6 +2257,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() @@ -2485,6 +2489,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) @@ -2520,6 +2525,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: @@ -3207,6 +3213,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) @@ -3428,6 +3435,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 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. @@ -3453,6 +3475,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()) @@ -3682,6 +3720,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. @@ -3789,6 +3828,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: @@ -3842,6 +3886,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 [] @@ -3937,6 +3982,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. @@ -4115,6 +4161,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: @@ -4168,6 +4219,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 @@ -4599,6 +4651,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. @@ -4691,6 +4744,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__( @@ -4716,6 +4774,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 757e2a15f8..3876f0e46a 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 @@ -2880,7 +2880,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: @@ -2904,8 +2904,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(idempotentHint=True) result = await server.call_tool("test_tool", param="test_value") @@ -2915,6 +2916,159 @@ 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_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(idempotentHint=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(readOnlyHint=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", + inputSchema={"type": "object", "properties": {}}, + annotations=types.ToolAnnotations(idempotentHint=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.""" @@ -3968,7 +4122,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 @@ -3987,6 +4149,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(readOnlyHint=True) + # Simulate connection failure by making call_tool raise ClosedResourceError once call_count = 0 @@ -7015,6 +7181,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 @@ -7029,7 +7198,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(idempotentHint=True) # Mock the session mock_session = MagicMock() @@ -7270,11 +7441,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(idempotentHint=True) with ( patch.object(tool, "connect", AsyncMock()) as mock_connect,