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

@@ -90,7 +90,11 @@
     InRunPendingToolTransitionOwner,
 )
 from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id
-from ._message_adapters import _APPROVAL_DECISION_IS_BOOLEAN_KEY, normalize_agui_input_messages
+from ._message_adapters import (
+    _APPROVAL_DECISION_IS_BOOLEAN_KEY,
+    _take_backend_snapshot_messages,
+    normalize_agui_input_messages,
+)
 from ._predictive_state import PredictiveStateHandler
 from ._tooling import collect_server_tools, merge_tools
 from ._run_common import (
@@ -2801,6 +2805,33 @@ 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],
+) -> None:
+    """Mark backend-restored replay messages for FIDES without importing the experimental module."""
+    if not messages:
+        return
+    for provider in getattr(agent, "context_providers", ()):
+        marker = getattr(provider, "_mark_replay_state_covered_messages", None)
+        if callable(marker):
+            marker(messages)
+
+
+def _fold_fides_uncovered_replay_messages(
+    agent: SupportsAgentRun,
+    session: AgentSession,
+    messages: Sequence[Message],
+    *,
+    complete_replay_messages: Sequence[Message] | None = None,
+) -> None:
+    """Fold uncovered replay into FIDES state before a snapshot-only return."""
+    for provider in getattr(agent, "context_providers", ()):
+        fold = getattr(provider, "_fold_uncovered_replay_messages", None)
+        if callable(fold):
+            fold(session, messages, complete_replay_messages=complete_replay_messages)
+
+
 def _legacy_tool_message_approval_resume(
     messages: Sequence[Mapping[str, Any]],
     *,
@@ -3004,6 +3035,7 @@ async def _run_agent_stream(
     )
 
     stored_snapshot = snapshot_session.stored
+    backend_snapshot_message_ids: set[int] = set()
     stored_pending_approval_interrupt_ids: set[str] = set()
     seeded_resume_from_snapshot = False
     if stored_snapshot is not None:
@@ -3073,6 +3105,8 @@ async def _run_agent_stream(
                 stored_interrupt=stored_snapshot.interrupt,
             )
             raw_messages = provider_suffix
+        if not config.use_service_session:
+            backend_snapshot_message_ids = {id(message) for message in raw_messages[: len(stored_snapshot.messages)]}
 
     # Initialize flow state with stored state plus request-provided overrides;
     # endpoint-deferred defaults apply only to keys missing from both.
@@ -3198,7 +3232,9 @@ async def _run_agent_stream(
         raw_messages,
         sanitize_tool_history=not (config.use_service_session and stored_snapshot is not None),
         protected_tool_call_ids=protected_tool_call_ids,
+        backend_snapshot_message_ids=backend_snapshot_message_ids,
     )
+    _mark_fides_state_covered_messages(agent, _take_backend_snapshot_messages(messages))
 
     if snapshot_seed_messages is not None:
         _, snapshot_messages = normalize_agui_input_messages(
@@ -3361,6 +3397,23 @@ async def save_thread_snapshot(
             ),
         )
 
+    def fold_fides_replay_before_snapshot_save() -> None:
+        complete_replay_snapshot = snapshot_messages
+        if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None:
+            complete_replay_snapshot = snapshot_session.resume_seeded_messages(complete_replay_snapshot)
+        complete_replay_messages, _ = normalize_agui_input_messages(
+            complete_replay_snapshot,
+            sanitize_tool_history=not (config.use_service_session and stored_snapshot is not None),
+            protected_tool_call_ids=protected_tool_call_ids,
+        )
+        _fold_fides_uncovered_replay_messages(
+            agent,
+            session,
+            messages,
+            complete_replay_messages=complete_replay_messages,
+        )
+        _mark_fides_state_covered_messages(agent, messages)
+
     authenticated_cancellations = [
         _approval_observer_response(occurrence, cancelled=True)
         for interrupt_id in cancelled_resume_ids
@@ -3613,6 +3666,7 @@ def merge_snapshot_interrupts(
         persisted_messages = snapshot_messages
         if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None:
             persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages)
+        fold_fides_replay_before_snapshot_save()
         await save_thread_snapshot(
             persisted_messages=persisted_messages,
             state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None,
@@ -3646,6 +3700,7 @@ def merge_snapshot_interrupts(
         if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None:
             persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages)
         preserved_interrupts = remaining_stored_interrupts(retired_approval_interrupt_ids)
+        fold_fides_replay_before_snapshot_save()
         await save_thread_snapshot(
             persisted_messages=persisted_messages,
             state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None,
@@ -3654,6 +3709,7 @@ def merge_snapshot_interrupts(
 
     # Handle confirm_changes response (state confirmation flow - emit confirmation and stop)
     if is_confirm_changes_response:
+        fold_fides_replay_before_snapshot_save()
         confirm_additional_properties = cast(dict[str, Any], messages[-1].additional_properties or {})
         confirm_interrupt_id = confirm_additional_properties.get("tool_call_id")
         confirm_remaining_interrupts = remaining_stored_interrupts(