Python: Harden FIDES security state and policy enforcement

docs/features/FIDES_IMPLEMENTATION_SUMMARY.md

@@ -4,6 +4,10 @@
 
 **FIDES**  is a comprehensive deterministic prompt injection defense system for the agent framework. The implementation provides label-based security mechanisms to defend against prompt injection attacks by tracking integrity and confidentiality of content throughout agent execution.
 
+FIDES enforces tool-result, hidden-variable, caller-replayed, and history-provider flows. Other `ContextProvider`
+implementations remain trusted plugins; provider-injected messages require an explicit provider/application trust
+boundary until a first-class provider-message labeling contract is defined.
+
 **🚀 Key Features:**
 - **Context Provider Pattern** - `SecureAgentConfig` extends `ContextProvider`, injecting tools, instructions, and middleware automatically
 - **Automatic Variable Hiding** - UNTRUSTED content is automatically hidden without requiring manual intervention
@@ -363,4 +367,8 @@ cd python/packages/core && ../../.venv/bin/pytest tests/test_security.py -v
 - **Full traceability**: Message-level label tracking
 - **Complete auditability**: All security events logged
 
-The system ensures that untrusted content never directly reaches the LLM context and that all tool calls are policy-checked based on the cumulative security state before execution.
+The system hides FIDES-labeled untrusted tool output from the LLM context and policy-checks tool calls against the
+cumulative security state before execution. Newly supplied assistant/tool messages are conservatively labeled even
+when state was restored. History-provider messages reconstruct missing state, while restored state remains
+authoritative for unstamped stored history; trusted non-history context providers remain responsible for the data they
+inject.

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py

@@ -2801,6 +2801,22 @@ def _split_service_session_input(
     return snapshot_messages[len(stored_snapshot_messages) :], snapshot_messages
 
 
+def _mark_fides_state_covered_messages(
+    agent: SupportsAgentRun,
+    messages: Sequence[Message],
+    *,
+    prefix_length: int,
+) -> None:
+    """Mark backend-restored replay messages for FIDES without importing the experimental module."""
+    if prefix_length <= 0:
+        return
+    covered_messages = messages[: min(prefix_length, len(messages))]
+    for provider in getattr(agent, "context_providers", ()):
+        marker = getattr(provider, "_mark_replay_state_covered_messages", None)
+        if callable(marker):
+            marker(covered_messages)
+
+
 def _legacy_tool_message_approval_resume(
     messages: Sequence[Mapping[str, Any]],
     *,
@@ -3004,6 +3020,7 @@ async def _run_agent_stream(
     )
 
     stored_snapshot = snapshot_session.stored
+    state_covered_prefix_length = 0
     stored_pending_approval_interrupt_ids: set[str] = set()
     seeded_resume_from_snapshot = False
     if stored_snapshot is not None:
@@ -3073,6 +3090,8 @@ async def _run_agent_stream(
                 stored_interrupt=stored_snapshot.interrupt,
             )
             raw_messages = provider_suffix
+        if not config.use_service_session:
+            state_covered_prefix_length = len(stored_snapshot.messages)
 
     # Initialize flow state with stored state plus request-provided overrides;
     # endpoint-deferred defaults apply only to keys missing from both.
@@ -3199,6 +3218,11 @@ async def _run_agent_stream(
         sanitize_tool_history=not (config.use_service_session and stored_snapshot is not None),
         protected_tool_call_ids=protected_tool_call_ids,
     )
+    _mark_fides_state_covered_messages(
+        agent,
+        messages,
+        prefix_length=state_covered_prefix_length,
+    )
 
     if snapshot_seed_messages is not None:
         _, snapshot_messages = normalize_agui_input_messages(

python/packages/core/agent_framework/_agents.py

@@ -1732,6 +1732,15 @@ async def _prepare_run_context(
                 providers=per_service_call_history_providers,
                 service_stores_history=service_stores_history,
             )
+        if per_service_call_history_middleware is not None:
+            existing_middleware = cast(
+                "MiddlewareTypes | Sequence[MiddlewareTypes] | None", effective_client_kwargs.get("middleware")
+            )
+            effective_client_kwargs["middleware"] = [
+                *_as_middleware_list(existing_middleware),
+                per_service_call_history_middleware,
+            ]
+
         provider_middleware = session_context.get_middleware()
         if provider_middleware:
             # Providers may only contribute chat/function middleware (enforced by
@@ -1751,15 +1760,6 @@ async def _prepare_run_context(
                     *provider_function_chat_middleware,
                 ]
 
-        if per_service_call_history_middleware is not None:
-            existing_middleware = cast(
-                "MiddlewareTypes | Sequence[MiddlewareTypes] | None", effective_client_kwargs.get("middleware")
-            )
-            effective_client_kwargs["middleware"] = [
-                *_as_middleware_list(existing_middleware),
-                per_service_call_history_middleware,
-            ]
-
         return {
             "session": active_session,
             "session_context": session_context,

python/packages/core/agent_framework/_mcp.py

@@ -193,6 +193,7 @@ class MCPSpecificApproval(TypedDict, total=False):
 _MCP_IS_TOOL_KEY = "_mcp_is_tool"
 _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name"
 _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY = "_mcp_tool_result_host_payload"
+_MCP_TOOL_RESULT_SECURITY_META_KEY = "_mcp_tool_result_security_meta"
 _MCP_PROGRESSIVE_LIST_TOOL_NAME = "list_mcp_tools"
 _MCP_PROGRESSIVE_LOAD_TOOL_NAME = "load_tool"
 _MCP_PROGRESSIVE_UNLOAD_TOOL_NAME = "unload_tool"
@@ -239,6 +240,7 @@ class MCPSpecificApproval(TypedDict, total=False):
 _MCP_SUBSCRIPTION_RETRY_INITIAL_DELAY = 0.1
 _MCP_SUBSCRIPTION_RETRY_MAX_DELAY = 2.0
 _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES = 1024 * 1024
+_MAX_MCP_SECURITY_META_SIZE_BYTES = 4096
 
 
 @dataclass
@@ -250,8 +252,45 @@ class _MCPHostPayloadCapture:
     apply_server_meta_to_model_items: bool
     host_payload: dict[str, Any] | None = None
     meta: dict[str, Any] | None = None
+    security_meta: dict[str, Any] | None = None
     recorded: bool = False
     meta_prepared: bool = False
+    security_meta_prepared: bool = False
+
+    def prepare_security_meta(self, mcp_type: Any) -> dict[str, Any] | None:
+        """Retain security-relevant MCP metadata outside generic payload budgets."""
+        if self.security_meta_prepared:
+            return self.security_meta
+        self.security_meta_prepared = True
+
+        raw_meta = getattr(mcp_type, "meta", None)
+        if not isinstance(raw_meta, Mapping) or "ifc" not in raw_meta:
+            return None
+
+        raw_meta_mapping = cast(Mapping[str, Any], raw_meta)
+        raw_security_meta: dict[str, Any] = {"ifc": raw_meta_mapping["ifc"]}
+        if _json_size_exceeds(raw_security_meta, _MAX_MCP_SECURITY_META_SIZE_BYTES):
+            logger.warning("MCP result _meta.ifc exceeds the independent security metadata limit.")
+            self.security_meta = {"retained": False}
+            return self.security_meta
+
+        try:
+            copied = cast(
+                dict[str, Any],
+                to_jsonable_python(raw_security_meta, fallback=str),
+            )
+        except (TypeError, ValueError):
+            logger.warning("MCP result _meta.ifc could not be retained safely.")
+            self.security_meta = {"retained": False}
+            return self.security_meta
+
+        if len(json.dumps(copied).encode("utf-8")) > _MAX_MCP_SECURITY_META_SIZE_BYTES:
+            logger.warning("MCP result _meta.ifc exceeds the independent security metadata limit.")
+            self.security_meta = {"retained": False}
+            return self.security_meta
+
+        self.security_meta = {"retained": True, "meta": copied}
+        return self.security_meta
 
     def prepare_meta(self, mcp_type: Any) -> dict[str, Any] | None:
         if self.meta_prepared:
@@ -281,6 +320,7 @@ def record(self, mcp_type: Any) -> None:
         if self.recorded:
             return
         self.recorded = True
+        self.prepare_security_meta(mcp_type)
         self.prepare_meta(mcp_type)
 
         effective_limit = self.max_size_bytes
@@ -314,7 +354,7 @@ def to_carrier(self) -> _FunctionResultCarrier:
             additional_properties=additional_properties,
             item_additional_properties=item_additional_properties,
             exclusive_outer_keys=frozenset({_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY}),
-            exclusive_item_keys=frozenset({"_meta"}),
+            exclusive_item_keys=frozenset({"_meta", _MCP_TOOL_RESULT_SECURITY_META_KEY}),
             result_already_parsed=True,
         )
 
@@ -326,8 +366,11 @@ def prepare_model_result(self, parsed: str | list[Content]) -> list[Content]:
             updated_item.additional_properties = dict(updated_item.additional_properties)
             updated_item.additional_properties.pop(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, None)
             updated_item.additional_properties.pop("_meta", None)
+            updated_item.additional_properties.pop(_MCP_TOOL_RESULT_SECURITY_META_KEY, None)
             if self.apply_server_meta_to_model_items and self.meta is not None:
                 updated_item.additional_properties["_meta"] = self.meta
+            if self.security_meta is not None:
+                updated_item.additional_properties[_MCP_TOOL_RESULT_SECURITY_META_KEY] = self.security_meta
             items[index] = updated_item
         return items
 
@@ -1340,6 +1383,7 @@ def _parse_tool_result_from_mcp(
 
         capture = _mcp_host_payload_capture.get()
         if capture is not None:
+            capture.prepare_security_meta(mcp_type)
             meta = capture.prepare_meta(mcp_type)
         else:
             raw_meta = mcp_type.meta

python/packages/core/agent_framework/_middleware.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import contextlib
+import contextvars
 import inspect
 import logging
 import sys
@@ -28,6 +29,11 @@
 
 logger = logging.getLogger(__name__)
 
+_ACTIVE_FUNCTION_MIDDLEWARE_PIPELINE: contextvars.ContextVar[tuple[Any, ...] | None] = contextvars.ContextVar(
+    "_active_function_middleware_pipeline",
+    default=None,
+)
+
 if sys.version_info >= (3, 13):
     from typing import TypeVar  # pragma: no cover
 else:
@@ -1573,7 +1579,11 @@ async def current_handler() -> None:
 
         first_handler = create_next_handler(0)
         # Don't suppress MiddlewareTermination - let it propagate to signal loop termination
-        await first_handler()
+        pipeline_token = _ACTIVE_FUNCTION_MIDDLEWARE_PIPELINE.set(tuple(self._middleware))
+        try:
+            await first_handler()
+        finally:
+            _ACTIVE_FUNCTION_MIDDLEWARE_PIPELINE.reset(pipeline_token)
 
         return context.result
 

python/packages/core/agent_framework/_tools.py

@@ -138,6 +138,9 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None)
 _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY: Final[str] = "_security_function_arguments"
 _PREPARED_ARGUMENTS_CONTEXT_KEY: Final[str] = "_prepared_function_arguments"
 _AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY: Final[str] = "_auto_prepare_function_arguments"
+_SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT: Final[
+    contextvars.ContextVar[tuple[Any, Callable[[Any], None]] | None]
+] = contextvars.ContextVar("_security_argument_materialization_context", default=None)
 _FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT: Final[str] = (
     "Function invocation limit reached before a final answer could be produced."
 )
@@ -1140,6 +1143,13 @@ async def invoke(
         from ._types import Content
         from .observability import OBSERVABILITY_SETTINGS
 
+        security_materialization = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.get()
+        if security_materialization is not None:
+            active_context, materialize = security_materialization
+            if context is not None and context is active_context:
+                materialize(context)
+                arguments = context.arguments
+
         configured_parser = self.result_parser
         skip_parsing = skip_parsing or configured_parser is SKIP_PARSING