Python: Apply MCP security labels to dynamically loaded tools

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**: Reduce response-stream hot-path overhead and avoid loading tool-approval middleware for agent runs without middleware ([#8971](https://github.com/microsoft/agent-framework/pull/8971))
 - **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))

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

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,40 +4453,78 @@ 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()
+        was_connected = self.is_connected
+        callback_bound = self._bind_function_load_callback()
+        try:
+            await self._mcp_tool.__aenter__()
+            await self._apply_labels()
+        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:
+                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()
+        was_connected = self.is_connected
+        callback_bound = self._bind_function_load_callback()
+        try:
+            await self._mcp_tool.connect()
+            await self._apply_labels()
+        except BaseException:
+            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:
+                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.")
-        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 --
 
@@ -4505,6 +4555,32 @@ def mcp_tool(self) -> Any:
 
     # -- Internal --
 
+    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 (
+            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,