From 9ce0aefc7876b6faa526a52716948fa981ab82bc Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 2 Oct 2026 13:28:35 +0200 Subject: [PATCH 1/5] Fix MCP security labels for dynamically loaded tools --- python/packages/core/agent_framework/_mcp.py | 18 ++ .../packages/core/agent_framework/security.py | 186 +++++++++------ python/packages/core/tests/test_security.py | 222 +++++++++++++++++- 3 files changed, 358 insertions(+), 68 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 242e94cae68..56f68cf3c3b 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -499,6 +499,7 @@ async def delete(self, *args: Any, **kwargs: Any) -> Any: # and returns (or awaits to) a truthy value to approve the request or a falsy # value to deny it. Both synchronous and asynchronous callables are supported. SamplingApprovalCallback = Callable[["types.CreateMessageRequestParams"], "bool | Coroutine[Any, Any, bool]"] +_MCPFunctionLoadCallback = Callable[[FunctionTool, Any], None] # region: Helpers @@ -1048,6 +1049,7 @@ def __init__( self._warn_sampling_deprecated(stacklevel=4) self._sampling_request_count = 0 self._functions: list[FunctionTool] = [] + self._function_load_callback: _MCPFunctionLoadCallback | None = None self.use_progressive_disclosure = use_progressive_disclosure self.always_load = always_load self._always_load_names = set(always_load or ()) @@ -2514,6 +2516,9 @@ async def _load_prompts_locked(self) -> None: params = types.PaginatedRequestParams(cursor=prompt_list.nextCursor) self._validate_config_names([*self._functions, *new_functions]) + if self._function_load_callback is not None: + for function in new_functions: + self._function_load_callback(function, None) self._functions.extend(new_functions) async def load_tools(self) -> None: @@ -2553,6 +2558,7 @@ async def _load_tools_locked(self) -> None: tool_call_meta_by_name: dict[str, dict[str, Any]] = {} tool_task_support_by_name: dict[str, str] = {} tool_param_names_by_name: dict[str, set[str]] = {} + tool_annotations_by_name: dict[str, Any] = {} params: types.PaginatedRequestParams | None = None while True: @@ -2588,6 +2594,7 @@ async def _load_tools_locked(self) -> None: raise ToolExecutionException("Failed to load tools.") for tool in tool_list.tools: + tool_annotations_by_name[tool.name] = tool.annotations if tool.meta is not None: tool_call_meta_by_name[tool.name] = _validate_mcp_meta(tool.meta) or {} @@ -2671,6 +2678,17 @@ async def _load_tools_locked(self) -> None: ] current_functions.extend(new_functions) self._validate_config_names(current_functions) + for function in current_functions: + properties = function.additional_properties or {} + if not properties.get(_MCP_IS_TOOL_KEY): + continue + remote_name = properties.get(_MCP_REMOTE_NAME_KEY) + if ( + isinstance(remote_name, str) + and remote_name in tool_annotations_by_name + and self._function_load_callback is not 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_task_support_by_name = tool_task_support_by_name diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 85f317b84da..46ef9d06d4c 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -4055,6 +4055,52 @@ def _map_mcp_annotations_to_labels( return (integrity, ConfidentialityLabel.PUBLIC, False) +def _apply_mcp_security_label_to_function( + function: FunctionTool, + annotations: Any, + *, + default_integrity: IntegrityLabel, + annotation_overrides: Mapping[str, tuple[IntegrityLabel, ConfidentialityLabel | None]] | None, + mark_write_tools_as_sinks: bool, + trust_server_ifc: bool, +) -> None: + """Apply one local MCP security policy to a remote function.""" + properties = function.additional_properties + if properties is None: + properties = {} + function.additional_properties = properties + remote_name = properties.get("_mcp_remote_name") + if not isinstance(remote_name, str): + return + + overrides = annotation_overrides or {} + if remote_name in overrides: + integrity, max_confidentiality = overrides[remote_name] + accepts_untrusted = False + else: + integrity, max_confidentiality, accepts_untrusted = _map_mcp_annotations_to_labels( + annotations, + default_integrity=default_integrity, + ) + + properties["source_integrity"] = integrity.value + if mark_write_tools_as_sinks and max_confidentiality is not None: + properties["max_allowed_confidentiality"] = max_confidentiality.value + else: + properties.pop("max_allowed_confidentiality", None) + properties["accepts_untrusted"] = accepts_untrusted + properties[_MCP_TRUST_SERVER_IFC_KEY] = trust_server_ifc + _wrap_mcp_function_for_ifc(function, default_integrity) + + logger.info( + "MCP auto-label: tool=%s integrity=%s max_confidentiality=%s accepts_untrusted=%s", + remote_name, + integrity.value, + max_confidentiality.value if max_confidentiality else "none", + accepts_untrusted, + ) + + @experimental(feature_id=ExperimentalFeature.FIDES) async def apply_mcp_security_labels( mcp_tool: Any, @@ -4072,6 +4118,7 @@ async def apply_mcp_security_labels( ``additional_properties``. The existing :class:`LabelTrackingFunctionMiddleware` picks these up automatically (Tier 2 label propagation), so **no middleware changes are needed**. + Currently hidden progressive-disclosure tools are included. Server annotations cannot relax local policy. Use ``annotation_overrides`` for explicit local per-tool static policy. Server result ``_meta.ifc`` is @@ -4080,6 +4127,8 @@ async def apply_mcp_security_labels( that result. ToolAnnotations remain non-authoritative in both modes. Call this **after** the ``MCPTool`` is connected (tools already loaded). + Use :class:`SecureMCPToolProxy` to keep the same policy bound to functions + discovered later in the connection lifecycle. Args: mcp_tool: A connected ``MCPTool`` instance (``MCPStdioTool``, @@ -4123,70 +4172,31 @@ async def apply_mcp_security_labels( if session is None: raise RuntimeError("MCPTool has no active session.") - # ------------------------------------------------------------------ - # 1. Fetch tool list (with annotations) from the server - # ------------------------------------------------------------------ from mcp import types as mcp_types - annotation_map: dict[str, Any] = {} # remote_name → ToolAnnotations | None + annotation_map: dict[str, Any] = {} params: mcp_types.PaginatedRequestParams | None = None while True: tool_list = await session.list_tools(params=params) - for t in tool_list.tools: - annotation_map[t.name] = t.annotations - if not tool_list or not tool_list.nextCursor: + for remote_tool in tool_list.tools: + annotation_map[remote_tool.name] = remote_tool.annotations + if not tool_list.nextCursor: break params = mcp_types.PaginatedRequestParams(cursor=tool_list.nextCursor) - # ------------------------------------------------------------------ - # 2. Patch each FunctionTool's additional_properties - # ------------------------------------------------------------------ - overrides = annotation_overrides or {} - functions: list[FunctionTool] = getattr(mcp_tool, "functions", []) - - for func in functions: - props = func.additional_properties - if props is None: - props = {} - func.additional_properties = props - - remote_name: str | None = props.get("_mcp_remote_name") - if remote_name is None: - continue - - # Check for explicit per-tool override first - if remote_name in overrides: - integrity, max_conf = overrides[remote_name] - accepts_untrusted = False # overrides must opt-in explicitly - else: - annotations = annotation_map.get(remote_name) - integrity, max_conf, accepts_untrusted = _map_mcp_annotations_to_labels( - annotations, default_integrity=default_integrity - ) - - # Patch source_integrity (Tier 2 - read by LabelTrackingFunctionMiddleware) - props["source_integrity"] = integrity.value - - # Patch sink constraint - if mark_write_tools_as_sinks and max_conf is not None: - props["max_allowed_confidentiality"] = max_conf.value - else: - props.pop("max_allowed_confidentiality", None) - - # Server annotations cannot authorize tainted input. - props["accepts_untrusted"] = accepts_untrusted - - # Local configuration controls result-label authority; MCP result - # metadata is attached to Content and cannot mutate tool properties. - props[_MCP_TRUST_SERVER_IFC_KEY] = trust_server_ifc - _wrap_mcp_function_for_ifc(func, default_integrity) - - logger.info( - "MCP auto-label: tool=%s integrity=%s max_confidentiality=%s accepts_untrusted=%s", - remote_name, - integrity.value, - max_conf.value if max_conf else "none", - accepts_untrusted, + loaded_functions = getattr(mcp_tool, "_functions", None) + if not isinstance(loaded_functions, list): + loaded_functions = getattr(mcp_tool, "functions", []) + for function in cast(list[FunctionTool], loaded_functions): + properties = function.additional_properties or {} + remote_name = properties.get("_mcp_remote_name") + _apply_mcp_security_label_to_function( + function, + annotation_map.get(remote_name) if isinstance(remote_name, str) else None, + default_integrity=default_integrity, + annotation_overrides=annotation_overrides, + mark_write_tools_as_sinks=mark_write_tools_as_sinks, + trust_server_ifc=trust_server_ifc, ) @@ -4318,7 +4328,9 @@ class SecureMCPToolProxy: Wraps any ``MCPTool`` subclass and calls :func:`apply_mcp_security_labels` automatically when entering the async - context manager (or when :meth:`connect` is called explicitly). + context manager (or when :meth:`connect` is called explicitly). The same + local policy is applied before later-discovered and progressively exposed + remote functions become callable. The proxy delegates ``functions``, ``is_connected``, and ``name`` to the wrapped tool. Pass ``proxy.tools`` (or ``proxy.functions``) directly to @@ -4441,39 +4453,56 @@ def __init__( self._annotation_overrides = annotation_overrides self._mark_write_tools_as_sinks = mark_write_tools_as_sinks self._trust_server_ifc = trust_server_ifc + self._function_load_callback = self._apply_function_labels # -- Async context manager -- async def __aenter__(self) -> SecureMCPToolProxy: """Enter context, connect the wrapped tool, and apply labels.""" - await self._mcp_tool.__aenter__() - await self._apply_labels() + self._bind_function_load_callback() + try: + await self._mcp_tool.__aenter__() + await self._apply_labels() + except BaseException: + self._unbind_function_load_callback() + raise return self async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: """Exit context and close the wrapped MCP tool.""" - await self._mcp_tool.__aexit__(exc_type, exc_val, exc_tb) + try: + await self._mcp_tool.__aexit__(exc_type, exc_val, exc_tb) + finally: + self._unbind_function_load_callback() # -- Explicit connect/disconnect -- async def connect(self) -> None: """Connect the underlying MCPTool and apply security labels.""" - await self._mcp_tool.connect() - await self._apply_labels() + self._bind_function_load_callback() + try: + await self._mcp_tool.connect() + await self._apply_labels() + except BaseException: + self._unbind_function_load_callback() + raise async def disconnect(self) -> None: """Disconnect the underlying MCPTool.""" - await self._mcp_tool.close() + try: + await self._mcp_tool.close() + finally: + self._unbind_function_load_callback() async def refresh_labels(self) -> None: - """Re-apply labels/wrappers for tools added while connected. + """Re-apply labels and wrappers to all currently discovered functions. - Some MCP servers can expose additional tools during a long-lived - connection. Call this to re-run annotation mapping and wrap newly - discovered tool callables without reconnecting. + Later-discovered functions are labeled automatically while the proxy + is active. """ if not self.is_connected: raise RuntimeError("MCPTool is not connected. Connect before refreshing labels.") + self._bind_function_load_callback() await self._apply_labels() # -- Delegated properties -- @@ -4505,6 +4534,29 @@ def mcp_tool(self) -> Any: # -- Internal -- + def _bind_function_load_callback(self) -> None: + current = self._mcp_tool._function_load_callback # pyright: ignore[reportPrivateUsage] + if current is not None and current is not self._function_load_callback: + raise RuntimeError("MCPTool is already wrapped by another SecureMCPToolProxy.") + self._mcp_tool._function_load_callback = self._function_load_callback # pyright: ignore[reportPrivateUsage] + + def _unbind_function_load_callback(self) -> None: + if ( + self._mcp_tool._function_load_callback # pyright: ignore[reportPrivateUsage] + is self._function_load_callback + ): + self._mcp_tool._function_load_callback = None # pyright: ignore[reportPrivateUsage] + + def _apply_function_labels(self, function: FunctionTool, annotations: Any) -> None: + _apply_mcp_security_label_to_function( + function, + annotations, + default_integrity=self._default_integrity, + annotation_overrides=self._annotation_overrides, + mark_write_tools_as_sinks=self._mark_write_tools_as_sinks, + trust_server_ifc=self._trust_server_ifc, + ) + async def _apply_labels(self) -> None: await apply_mcp_security_labels( self._mcp_tool, diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index b2c3aaf82e8..db08688b410 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -6,10 +6,11 @@ import json import logging import math +import warnings from datetime import timedelta from types import MappingProxyType, SimpleNamespace from typing import Any, cast -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest from pydantic import BaseModel, field_validator @@ -6090,6 +6091,59 @@ async def fake_call(**kwargs: Any) -> list[Content]: return mcp_tool, function +def _make_mcp_tool_definition(name: str, *, open_world: bool = False) -> Any: + from mcp import types as mcp_types + + return mcp_types.Tool( + name=name, + description=f"{name} description", + inputSchema={"type": "object", "properties": {}}, + annotations=mcp_types.ToolAnnotations(readOnlyHint=False, openWorldHint=open_world), + ) + + +def _make_connected_mcp_discovery_tool( + *, + progressive: bool = False, + always_load: list[str] | None = None, + result_meta: dict[str, Any] | None = None, +) -> Any: + from mcp import types as mcp_types + + from agent_framework._mcp import MCPTool + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + mcp_tool = MCPTool( # type: ignore[abstract] + name="helper", + load_prompts=False, + use_progressive_disclosure=progressive, + always_load=always_load, + ) + mcp_tool.is_connected = True + mcp_tool.session = AsyncMock() + mcp_tool.session.call_tool = AsyncMock( + return_value=mcp_types.CallToolResult( + content=[mcp_types.TextContent(type="text", text="payload")], + _meta=result_meta or {"ifc": {"integrity": "trusted", "confidentiality": "private"}}, + ) + ) + return mcp_tool + + +def _find_mcp_function(mcp_tool: Any, remote_name: str) -> FunctionTool: + return next( + function + for function in mcp_tool._functions + if (function.additional_properties or {}).get("_mcp_remote_name") == remote_name + ) + + +async def _invoke_mcp_function(function: FunctionTool) -> Any: + context = FunctionInvocationContext(function=function, arguments={}) + return await function.invoke(arguments={}, context=context, skip_parsing=True) + + # --------------------------------------------------------------------------- # IFC labels from MCP _meta payload # --------------------------------------------------------------------------- @@ -6668,6 +6722,172 @@ async def test_secure_mcp_proxy_configures_result_authority(self, trust_server_i ) assert result[0].additional_properties["security_label"] == expected_label + async def test_secure_mcp_proxy_labels_notification_reload_before_publication(self): + from mcp import types as mcp_types + + from agent_framework.security import SecureMCPToolProxy + + initial_tool = _make_mcp_tool_definition("initial_sink", open_world=False) + late_tool = _make_mcp_tool_definition("late_sink", open_world=True) + mcp_tool = _make_connected_mcp_discovery_tool() + mcp_tool.session.list_tools = AsyncMock(return_value=mcp_types.ListToolsResult(tools=[initial_tool])) + await mcp_tool.load_tools() + + proxy = SecureMCPToolProxy(mcp_tool, default_integrity=IntegrityLabel.TRUSTED) + await proxy.refresh_labels() + initial_function = _find_mcp_function(mcp_tool, "initial_sink") + assert initial_function.additional_properties is not None + assert initial_function.additional_properties["source_integrity"] == "trusted" + + reloaded_initial_tool = _make_mcp_tool_definition("initial_sink", open_world=True) + mcp_tool.session.list_tools.return_value = mcp_types.ListToolsResult(tools=[reloaded_initial_tool, late_tool]) + notification = Mock(spec=mcp_types.ServerNotification) + notification.root = Mock() + notification.root.method = "notifications/tools/list_changed" + + await mcp_tool.message_handler(notification) + pending_reloads = list(mcp_tool._pending_reload_tasks) + assert pending_reloads + await asyncio.gather(*pending_reloads) + + assert _find_mcp_function(mcp_tool, "initial_sink") is initial_function + assert initial_function.additional_properties is not None + assert initial_function.additional_properties["source_integrity"] == "untrusted" + late_function = _find_mcp_function(mcp_tool, "late_sink") + assert late_function.additional_properties is not None + assert late_function.additional_properties["source_integrity"] == "untrusted" + assert late_function.additional_properties["max_allowed_confidentiality"] == "public" + assert late_function.additional_properties["accepts_untrusted"] is False + assert late_function.additional_properties["_mcp_trust_server_ifc"] is False + assert getattr(late_function.func, "_ifc_wrapped", False) is True + + context = FunctionInvocationContext(function=late_function, arguments={}) + context.metadata["context_label"] = ContentLabel(confidentiality=ConfidentialityLabel.PRIVATE) + + async def execute() -> None: + pytest.fail("The PRIVATE call should be blocked before invoking the MCP server.") + + with pytest.raises(MiddlewareTermination): + await PolicyEnforcementFunctionMiddleware(block_on_violation=True).process(context, execute) + + assert "exfiltration" in context.result["error"].lower() + mcp_tool.session.call_tool.assert_not_called() + + result = await _invoke_mcp_function(late_function) + assert result[0].additional_properties["security_label"] == { + "integrity": "untrusted", + "confidentiality": "private", + } + assert "_meta" not in result[0].additional_properties + + async def test_secure_mcp_proxy_labels_progressive_hidden_function_before_exposure(self): + from mcp import types as mcp_types + + from agent_framework.security import SecureMCPToolProxy + + initial_tool = _make_mcp_tool_definition("initial_sink") + hidden_tool = _make_mcp_tool_definition("hidden_sink") + mcp_tool = _make_connected_mcp_discovery_tool( + progressive=True, + always_load=["initial_sink"], + ) + mcp_tool.session.list_tools = AsyncMock( + return_value=mcp_types.ListToolsResult(tools=[initial_tool, hidden_tool]) + ) + proxy = SecureMCPToolProxy(mcp_tool) + await proxy.refresh_labels() + await mcp_tool.load_tools() + + hidden_function = _find_mcp_function(mcp_tool, "hidden_sink") + assert hidden_function not in mcp_tool.functions + assert hidden_function.additional_properties is not None + assert hidden_function.additional_properties["source_integrity"] == "untrusted" + assert hidden_function.additional_properties["max_allowed_confidentiality"] == "public" + assert getattr(hidden_function.func, "_ifc_wrapped", False) is True + + load_function = next(function for function in mcp_tool.functions if function.name == "load_tool") + context = FunctionInvocationContext( + function=load_function, + arguments={"tool": "hidden_sink"}, + tools=list(mcp_tool.functions), + ) + await load_function.invoke(arguments={"tool": "hidden_sink"}, context=context) + + assert context.tools is not None + assert hidden_function in context.tools + policy_context = FunctionInvocationContext(function=hidden_function, arguments={}) + policy_context.metadata["context_label"] = ContentLabel(confidentiality=ConfidentialityLabel.PRIVATE) + + async def execute() -> None: + pytest.fail("The PRIVATE call should be blocked before invoking the MCP server.") + + with pytest.raises(MiddlewareTermination): + await PolicyEnforcementFunctionMiddleware(block_on_violation=True).process(policy_context, execute) + + assert "exfiltration" in policy_context.result["error"].lower() + mcp_tool.session.call_tool.assert_not_called() + + result = await _invoke_mcp_function(hidden_function) + assert result[0].additional_properties["security_label"] == { + "integrity": "untrusted", + "confidentiality": "private", + } + assert "_meta" not in result[0].additional_properties + + async def test_secure_mcp_proxy_applies_local_policy_to_later_tool(self): + from mcp import types as mcp_types + + from agent_framework.security import SecureMCPToolProxy + + later_tool = _make_mcp_tool_definition("later_tool", open_world=True) + mcp_tool = _make_connected_mcp_discovery_tool( + result_meta={"ifc": {"integrity": "trusted", "confidentiality": "public"}} + ) + mcp_tool.session.list_tools = AsyncMock(return_value=mcp_types.ListToolsResult(tools=[later_tool])) + proxy = SecureMCPToolProxy( + mcp_tool, + annotation_overrides={"later_tool": (IntegrityLabel.TRUSTED, ConfidentialityLabel.PRIVATE)}, + trust_server_ifc=True, + ) + await proxy.refresh_labels() + await mcp_tool.load_tools() + + function = _find_mcp_function(mcp_tool, "later_tool") + assert function.additional_properties is not None + assert function.additional_properties["source_integrity"] == "trusted" + assert function.additional_properties["max_allowed_confidentiality"] == "private" + assert function.additional_properties["_mcp_trust_server_ifc"] is True + result = await _invoke_mcp_function(function) + assert result[0].additional_properties["security_label"] == { + "integrity": "trusted", + "confidentiality": "public", + } + + await proxy.disconnect() + assert mcp_tool._function_load_callback is None + + async def test_plain_mcp_tool_remains_unlabeled(self): + from mcp import types as mcp_types + + plain_tool = _make_mcp_tool_definition("plain_tool") + mcp_tool = _make_connected_mcp_discovery_tool() + mcp_tool.session.list_tools = AsyncMock(return_value=mcp_types.ListToolsResult(tools=[plain_tool])) + + await mcp_tool.load_tools() + + function = _find_mcp_function(mcp_tool, "plain_tool") + assert function.additional_properties == { + "_mcp_remote_name": "plain_tool", + "_mcp_normalized_name": "plain_tool", + "_mcp_is_tool": True, + } + assert getattr(function.func, "_ifc_wrapped", False) is False + result = await _invoke_mcp_function(function) + assert result[0].additional_properties["_meta"] == { + "ifc": {"integrity": "trusted", "confidentiality": "private"} + } + assert "security_label" not in result[0].additional_properties + async def test_wrap_mcp_function_str_result_passes_through(self): """``str`` results (no per-item containers) are not modified by the wrapper.""" from agent_framework.security import _wrap_mcp_function_for_ifc From 5a0a7714d006aec14d25b9eb03ca8c0fc8992795 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 2 Oct 2026 13:36:22 +0200 Subject: [PATCH 2/5] Add MCP security label fix to 1.20.0 changelog --- python/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 33c1b85262a..e78e653ec2d 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -36,6 +36,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **docs**, **samples**: Cross-reference security patterns, clarify file-search retrieval scope, improve local-shell filtering guidance, and correct README instructions ([#8726](https://github.com/microsoft/agent-framework/pull/8726), [#8591](https://github.com/microsoft/agent-framework/pull/8591), [#8766](https://github.com/microsoft/agent-framework/pull/8766), [#8822](https://github.com/microsoft/agent-framework/pull/8822)) ### Fixed +- **agent-framework-core**: Apply MCP security labels before dynamically discovered and progressively disclosed tools become callable ([#8972](https://github.com/microsoft/agent-framework/pull/8972)) - **agent-framework-core**: Preserve functional-workflow replay identity, scope invocation arguments and activity IDs to workflow runs, deduplicate fan-out targets, propagate nested cancellation arguments, and support orchestration checkpoint deserialization ([#8887](https://github.com/microsoft/agent-framework/pull/8887), [#8596](https://github.com/microsoft/agent-framework/pull/8596), [#8549](https://github.com/microsoft/agent-framework/pull/8549), [#8739](https://github.com/microsoft/agent-framework/pull/8739), [#8602](https://github.com/microsoft/agent-framework/pull/8602), [#8723](https://github.com/microsoft/agent-framework/pull/8723)) - **agent-framework-core**: Preserve repeated history turns, compaction through middleware rewrites, checkpoint dictionary subclasses, response metadata, annotations, native vector keys, integer precision, and boolean set values ([#8800](https://github.com/microsoft/agent-framework/pull/8800), [#8671](https://github.com/microsoft/agent-framework/pull/8671), [#8519](https://github.com/microsoft/agent-framework/pull/8519), [#8464](https://github.com/microsoft/agent-framework/pull/8464), [#8701](https://github.com/microsoft/agent-framework/pull/8701), [#8535](https://github.com/microsoft/agent-framework/pull/8535), [#8637](https://github.com/microsoft/agent-framework/pull/8637), [#8681](https://github.com/microsoft/agent-framework/pull/8681)) - **agent-framework-core**: Improve MCP skill archive and resource validation, stop failed MCP lifecycle owners, distinguish missing tool arguments from parse failures, settle fatal unknown calls, and omit absent MCP prompt-argument descriptions ([#8937](https://github.com/microsoft/agent-framework/pull/8937), [#8690](https://github.com/microsoft/agent-framework/pull/8690), [#8755](https://github.com/microsoft/agent-framework/pull/8755), [#8609](https://github.com/microsoft/agent-framework/pull/8609), [#8640](https://github.com/microsoft/agent-framework/pull/8640), [#8733](https://github.com/microsoft/agent-framework/pull/8733)) From e7b5761cd48cad36fe28d7da64bf7e8ce45c973c Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 2 Oct 2026 13:40:21 +0200 Subject: [PATCH 3/5] Preserve MCP security binding on setup failures --- .../packages/core/agent_framework/security.py | 40 ++++++++--- python/packages/core/tests/test_security.py | 69 +++++++++++++++++++ 2 files changed, 101 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 46ef9d06d4c..f496c4544c3 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -4459,12 +4459,20 @@ def __init__( async def __aenter__(self) -> SecureMCPToolProxy: """Enter context, connect the wrapped tool, and apply labels.""" - self._bind_function_load_callback() + was_connected = self.is_connected + callback_bound = self._bind_function_load_callback() try: await self._mcp_tool.__aenter__() await self._apply_labels() - except BaseException: - self._unbind_function_load_callback() + except BaseException as ex: + if not was_connected and self.is_connected: + try: + await self._mcp_tool.__aexit__(type(ex), ex, ex.__traceback__) + finally: + if callback_bound: + self._unbind_function_load_callback() + elif callback_bound and not self.is_connected: + self._unbind_function_load_callback() raise return self @@ -4479,12 +4487,20 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: async def connect(self) -> None: """Connect the underlying MCPTool and apply security labels.""" - self._bind_function_load_callback() + was_connected = self.is_connected + callback_bound = self._bind_function_load_callback() try: await self._mcp_tool.connect() await self._apply_labels() except BaseException: - self._unbind_function_load_callback() + if not was_connected and self.is_connected: + try: + await self._mcp_tool.close() + finally: + if callback_bound: + self._unbind_function_load_callback() + elif callback_bound and not self.is_connected: + self._unbind_function_load_callback() raise async def disconnect(self) -> None: @@ -4502,8 +4518,13 @@ async def refresh_labels(self) -> None: """ if not self.is_connected: raise RuntimeError("MCPTool is not connected. Connect before refreshing labels.") - self._bind_function_load_callback() - await self._apply_labels() + callback_bound = self._bind_function_load_callback() + try: + await self._apply_labels() + except BaseException: + if callback_bound: + self._unbind_function_load_callback() + raise # -- Delegated properties -- @@ -4534,11 +4555,14 @@ def mcp_tool(self) -> Any: # -- Internal -- - def _bind_function_load_callback(self) -> None: + def _bind_function_load_callback(self) -> bool: current = self._mcp_tool._function_load_callback # pyright: ignore[reportPrivateUsage] + if current is self._function_load_callback: + return False if current is not None and current is not self._function_load_callback: raise RuntimeError("MCPTool is already wrapped by another SecureMCPToolProxy.") self._mcp_tool._function_load_callback = self._function_load_callback # pyright: ignore[reportPrivateUsage] + return True def _unbind_function_load_callback(self) -> None: if ( diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 80d0967fd77..1d4c5c099a8 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -6866,6 +6866,75 @@ async def test_secure_mcp_proxy_applies_local_policy_to_later_tool(self): await proxy.disconnect() assert mcp_tool._function_load_callback is None + @pytest.mark.parametrize("entrypoint", ["connect", "enter"]) + @pytest.mark.parametrize("failure_type", [RuntimeError, asyncio.CancelledError]) + @pytest.mark.parametrize("already_connected", [False, True]) + async def test_secure_mcp_proxy_setup_failure_preserves_live_connection_binding( + self, + entrypoint: str, + failure_type: type[BaseException], + already_connected: bool, + ): + from agent_framework.security import SecureMCPToolProxy + + mcp_tool = _make_connected_mcp_discovery_tool() + mcp_tool.is_connected = already_connected + proxy = SecureMCPToolProxy(mcp_tool) + if already_connected: + mcp_tool._function_load_callback = proxy._function_load_callback + + async def open_connection(*_args: Any) -> Any: + mcp_tool.is_connected = True + return mcp_tool + + async def close_connection(*_args: Any) -> None: + mcp_tool.is_connected = False + + open_mock = AsyncMock(side_effect=open_connection) + close_mock = AsyncMock(side_effect=close_connection) + if entrypoint == "connect": + mcp_tool.connect = open_mock # type: ignore[method-assign] + mcp_tool.close = close_mock # type: ignore[method-assign] + operation = proxy.connect + else: + mcp_tool.__aenter__ = open_mock # type: ignore[method-assign] + mcp_tool.__aexit__ = close_mock # type: ignore[method-assign] + operation = proxy.__aenter__ + proxy._apply_labels = AsyncMock(side_effect=failure_type("label refresh failed")) # type: ignore[method-assign] + + with pytest.raises(failure_type): + await operation() + + if already_connected: + assert mcp_tool.is_connected is True + assert mcp_tool._function_load_callback is proxy._function_load_callback + close_mock.assert_not_called() + else: + assert mcp_tool.is_connected is False + assert mcp_tool._function_load_callback is None + close_mock.assert_awaited_once() + + @pytest.mark.parametrize("failure_type", [RuntimeError, asyncio.CancelledError]) + @pytest.mark.parametrize("already_bound", [False, True]) + async def test_secure_mcp_proxy_refresh_failure_only_removes_new_binding( + self, + failure_type: type[BaseException], + already_bound: bool, + ): + from agent_framework.security import SecureMCPToolProxy + + mcp_tool = _make_connected_mcp_discovery_tool() + proxy = SecureMCPToolProxy(mcp_tool) + if already_bound: + mcp_tool._function_load_callback = proxy._function_load_callback + proxy._apply_labels = AsyncMock(side_effect=failure_type("label refresh failed")) # type: ignore[method-assign] + + with pytest.raises(failure_type): + await proxy.refresh_labels() + + expected_callback = proxy._function_load_callback if already_bound else None + assert mcp_tool._function_load_callback is expected_callback + async def test_plain_mcp_tool_remains_unlabeled(self): from mcp import types as mcp_types From a08adc8efe4dc7c3c0d568b0dcbc3fd8e62b5106 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 2 Oct 2026 13:47:25 +0200 Subject: [PATCH 4/5] Fix MCP proxy failure test typing --- python/packages/core/tests/test_security.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 1d4c5c099a8..8b9446d502b 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -6112,9 +6112,13 @@ def _make_connected_mcp_discovery_tool( from agent_framework._mcp import MCPTool + class _ConcreteMCPTool(MCPTool): + def get_mcp_client(self): + raise NotImplementedError + with warnings.catch_warnings(): warnings.simplefilter("ignore") - mcp_tool = MCPTool( # type: ignore[abstract] + mcp_tool = _ConcreteMCPTool( name="helper", load_prompts=False, use_progressive_disclosure=progressive, @@ -6895,15 +6899,16 @@ async def close_connection(*_args: Any) -> None: if entrypoint == "connect": mcp_tool.connect = open_mock # type: ignore[method-assign] mcp_tool.close = close_mock # type: ignore[method-assign] - operation = proxy.connect else: mcp_tool.__aenter__ = open_mock # type: ignore[method-assign] mcp_tool.__aexit__ = close_mock # type: ignore[method-assign] - operation = proxy.__aenter__ proxy._apply_labels = AsyncMock(side_effect=failure_type("label refresh failed")) # type: ignore[method-assign] with pytest.raises(failure_type): - await operation() + if entrypoint == "connect": + await proxy.connect() + else: + await proxy.__aenter__() if already_connected: assert mcp_tool.is_connected is True From 6e3d2e75f2a9ffd546fd1cb5e8839dc5beacec16 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 2 Oct 2026 13:55:10 +0200 Subject: [PATCH 5/5] Roll back newly acquired MCP security bindings --- python/packages/core/agent_framework/security.py | 4 ++-- python/packages/core/tests/test_security.py | 13 +++++++++---- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index f496c4544c3..0a7815c699e 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -4471,7 +4471,7 @@ async def __aenter__(self) -> SecureMCPToolProxy: finally: if callback_bound: self._unbind_function_load_callback() - elif callback_bound and not self.is_connected: + elif callback_bound: self._unbind_function_load_callback() raise return self @@ -4499,7 +4499,7 @@ async def connect(self) -> None: finally: if callback_bound: self._unbind_function_load_callback() - elif callback_bound and not self.is_connected: + elif callback_bound: self._unbind_function_load_callback() raise diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 8b9446d502b..1003aca53b8 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -6872,19 +6872,24 @@ async def test_secure_mcp_proxy_applies_local_policy_to_later_tool(self): @pytest.mark.parametrize("entrypoint", ["connect", "enter"]) @pytest.mark.parametrize("failure_type", [RuntimeError, asyncio.CancelledError]) - @pytest.mark.parametrize("already_connected", [False, True]) + @pytest.mark.parametrize( + ("already_connected", "already_bound"), + [(False, False), (True, False), (True, True)], + ids=["fresh", "connected-unbound", "connected-bound"], + ) async def test_secure_mcp_proxy_setup_failure_preserves_live_connection_binding( self, entrypoint: str, failure_type: type[BaseException], already_connected: bool, + already_bound: bool, ): from agent_framework.security import SecureMCPToolProxy mcp_tool = _make_connected_mcp_discovery_tool() mcp_tool.is_connected = already_connected proxy = SecureMCPToolProxy(mcp_tool) - if already_connected: + if already_bound: mcp_tool._function_load_callback = proxy._function_load_callback async def open_connection(*_args: Any) -> Any: @@ -6912,12 +6917,12 @@ async def close_connection(*_args: Any) -> None: if already_connected: assert mcp_tool.is_connected is True - assert mcp_tool._function_load_callback is proxy._function_load_callback close_mock.assert_not_called() else: assert mcp_tool.is_connected is False - assert mcp_tool._function_load_callback is None close_mock.assert_awaited_once() + expected_callback = proxy._function_load_callback if already_bound else None + assert mcp_tool._function_load_callback is expected_callback @pytest.mark.parametrize("failure_type", [RuntimeError, asyncio.CancelledError]) @pytest.mark.parametrize("already_bound", [False, True])