Python: Harden FIDES replay and approval authority

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
@@ -355,12 +359,16 @@ cd python/packages/core && ../../.venv/bin/pytest tests/test_security.py -v
 
 **FIDES** provides a comprehensive, deterministic defense against prompt injection attacks with:
 
-- **Zero-effort protection**: Automatic variable hiding for developers
+- **Automatic result protection**: FIDES-labeled untrusted tool output is hidden automatically
 - **Context provider pattern**: `SecureAgentConfig` extends `ContextProvider` for automatic setup
 - **Granular control**: Per-item embedded labels via `Content.from_text()` for mixed-trust data
 - **Easy configuration**: `SecureAgentConfig` for one-line setup
 - **Data safety**: Exfiltration prevention via confidentiality gates
 - **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. Stored provider history is reconstructed conservatively whenever it is loaded; FIDES does not
+infer trust through transcript persistence, deduplication, compaction, or merging. Trusted non-history context providers
+remain responsible for adding restriction labels or hiding the data they inject.

python/packages/ag-ui/tests/ag_ui/test_endpoint.py

@@ -7358,10 +7358,10 @@ async def test_endpoint_fides_replacement_rotates_lifecycle_generation(
     assert executed == ["guarded"]
 
 
-async def test_endpoint_fides_approval_uses_lifecycle_bound_request_generation(
+async def test_endpoint_fides_approval_without_replay_provenance_fails_closed(
     streaming_chat_client_stub: Any,
 ) -> None:
-    """AG-UI preserves trusted request generation while rebinding to occurrence identity."""
+    """AG-UI snapshot replay requires explicit provenance before FIDES can reuse one approval."""
     client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint(
         streaming_chat_client_stub
     )
@@ -7389,11 +7389,16 @@ async def test_endpoint_fides_approval_uses_lifecycle_bound_request_generation(
 
     assert approved.status_code == 200
     results = [event for event in _decode_sse_events(approved) if event.get("type") == "TOOL_CALL_RESULT"]
-    assert [(event["toolCallId"], event["content"]) for event in results] == [("call-guarded", "guarded result")]
-    assert executed == ["guarded"]
+    assert results == []
+    assert executed == []
+    finished = [event for event in _decode_sse_events(approved) if event.get("type") == "RUN_FINISHED"]
+    replacement_interrupts = _run_finished_interrupts(finished[-1])
+    assert len(replacement_interrupts) == 1
+    assert replacement_interrupts[0]["id"] != approval_id
+    assert replacement_interrupts[0]["toolCallId"] == "call-guarded"
     snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id)
     assert snapshot is not None and snapshot.session_state is not None
-    assert snapshot.session_state[security.source_id]["pending_policy_approvals"] == {}
+    assert snapshot.session_state[security.source_id]["pending_policy_approvals"]
 
 
 async def test_endpoint_agent_approval_resume_persists_replayable_tool_results(streaming_chat_client_stub):

python/packages/core/agent_framework/_agents.py

@@ -1641,6 +1641,7 @@ async def _prepare_run_context(
             session=active_session,
             input_messages=input_messages,
             options=opts,
+            service_stores_history=bool(service_stores_history),
         )
         default_additional_args = chat_options.pop("additional_function_arguments", None)
         if isinstance(default_additional_args, Mapping):
@@ -1756,8 +1757,19 @@ async def _prepare_run_context(
                 agent=self,
                 session=active_session,
                 providers=per_service_call_history_providers,
+                context_providers=self.context_providers,
+                run_context=session_context,
                 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
@@ -1777,15 +1789,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,
-            ]
-
         evaluation_tools: dict[str, dict[str, Any]] = {}
         if isinstance(self.client, BaseChatClient):
             effective_client_kwargs["_evaluation_tool_capture"] = partial(_capture_evaluation_tools, evaluation_tools)
@@ -1813,6 +1816,7 @@ async def _prepare_session_and_messages(
         session: AgentSession | None,
         input_messages: list[Message] | None = None,
         options: dict[str, Any] | None = None,
+        service_stores_history: bool = False,
     ) -> tuple[SessionContext, dict[str, Any]]:
         """Prepare the session context and messages for agent execution.
 
@@ -1823,6 +1827,7 @@ async def _prepare_session_and_messages(
             session: The conversation session (None for stateless invocation).
             input_messages: Messages to process.
             options: Runtime options dict (already copied, safe to mutate).
+            service_stores_history: Whether the service owns history for this run.
 
         Returns:
             A tuple containing:
@@ -1874,6 +1879,40 @@ async def _prepare_session_and_messages(
                 state=provider_session.state.setdefault(provider.source_id, {}),
             )
 
+        loaded_context = session_context
+        requires_loaded_context = any(
+            provider.requires_loaded_context_before_invocation for provider in self.context_providers
+        )
+        if per_service_call_history_required and not service_stores_history and requires_loaded_context:
+            if provider_session is None:
+                raise RuntimeError("Provider session must be available when per-service history is classified.")
+            loaded_context = SessionContext(
+                session_id=provider_session.session_id,
+                service_session_id=provider_session.service_session_id,
+                input_messages=list(session_context.input_messages),
+                options=dict(session_context.options),
+                metadata=session_context.metadata,
+            )
+            for provider in self.context_providers:
+                if isinstance(provider, HistoryProvider) and provider.load_messages:
+                    await provider.before_run(
+                        agent=self,
+                        session=provider_session,
+                        context=loaded_context,
+                        state=provider_session.state.setdefault(provider.source_id, {}),
+                    )
+
+        if self.context_providers:
+            if provider_session is None:
+                raise RuntimeError("Provider session must be available when loaded context is finalized.")
+            for provider in self.context_providers:
+                await provider.after_context_loaded(
+                    agent=self,
+                    session=provider_session,
+                    context=loaded_context,
+                    state=provider_session.state.setdefault(provider.source_id, {}),
+                )
+
         # Merge provider-contributed tools into chat_options
         if session_context.tools:
             if chat_options.get("tools") is not None:

python/packages/core/agent_framework/_sessions.py

@@ -777,6 +777,7 @@ class ContextProvider:
     """
 
     after_run_once_per_turn: bool = False
+    requires_loaded_context_before_invocation: bool = False
 
     def __init__(self, source_id: str):
         """Initialize the provider.
@@ -828,6 +829,16 @@ async def after_run(
                 Full cross-provider state remains available at ``session.state``.
         """
 
+    async def after_context_loaded(
+        self,
+        *,
+        agent: SupportsAgentRun,
+        session: AgentSession,
+        context: SessionContext,
+        state: dict[str, Any],
+    ) -> None:
+        """Finalize context after every loadable provider has run."""
+
 
 def _approval_controls_to_keep(messages: Sequence[Message]) -> set[int]:
     unresolved_requests_by_id: dict[str, Content] = {}
@@ -1089,6 +1100,23 @@ def _get_context_messages_to_store(self, context: SessionContext) -> list[Messag
             return context.get_messages(sources=self.store_context_from)
         return context.get_messages(exclude_sources={self.source_id})
 
+    def _get_messages_to_store(self, context: SessionContext, session: AgentSession) -> list[Message]:
+        """Return the exact framework-owned messages this provider will persist."""
+        messages_to_store: list[Message] = []
+        messages_to_store.extend(self._get_context_messages_to_store(context))
+        if self.store_inputs:
+            messages_to_store.extend(context.input_messages)
+        if self.store_outputs and context.response and context.response.messages:
+            output_messages = context.response.messages
+            if (
+                self.load_messages
+                and self.store_inputs
+                and session.state.get(_WORKFLOW_DEFER_COMPUTER_FUNCTION_RESULTS_KEY)
+            ):
+                output_messages = _without_deferred_workflow_function_results(output_messages)
+            messages_to_store.extend(output_messages)
+        return messages_to_store
+
     async def before_run(
         self,
         *,
@@ -1110,19 +1138,7 @@ async def after_run(
         state: dict[str, Any],
     ) -> None:
         """Store messages based on configuration."""
-        messages_to_store: list[Message] = []
-        messages_to_store.extend(self._get_context_messages_to_store(context))
-        if self.store_inputs:
-            messages_to_store.extend(context.input_messages)
-        if self.store_outputs and context.response and context.response.messages:
-            output_messages = context.response.messages
-            if (
-                self.load_messages
-                and self.store_inputs
-                and session.state.get(_WORKFLOW_DEFER_COMPUTER_FUNCTION_RESULTS_KEY)
-            ):
-                output_messages = _without_deferred_workflow_function_results(output_messages)
-            messages_to_store.extend(output_messages)
+        messages_to_store = self._get_messages_to_store(context, session)
         if messages_to_store:
             await self.save_messages(context.session_id, messages_to_store, state=state)
 
@@ -1601,6 +1617,8 @@ def __init__(
         agent: SupportsAgentRun,
         session: AgentSession,
         providers: Sequence[HistoryProvider],
+        context_providers: Sequence[ContextProvider],
+        run_context: SessionContext,
         service_stores_history: bool = False,
     ) -> None:
         """Initialize the middleware.
@@ -1609,6 +1627,8 @@ def __init__(
             agent: The agent that owns the history providers.
             session: The active session for the current run.
             providers: The history providers participating in per-service-call persistence.
+            context_providers: All providers participating in the owning run.
+            run_context: The owning run's SessionContext.
             service_stores_history: When True, the chat client stores history server-side. The
                 middleware then skips loading providers and leaves the real conversation id
                 untouched, persisting each service call without driving the function loop with a
@@ -1618,6 +1638,8 @@ def __init__(
         self._agent = agent
         self._session = session
         self._providers = list(providers)
+        self._context_providers = list(context_providers)
+        self._run_context = run_context
         self._service_stores_history = service_stores_history
         self._merged_provider_tool_names: set[str] = set()
         self._merged_provider_instructions: set[str] = set()
@@ -1629,6 +1651,7 @@ async def _prepare_service_call_context(self, messages: Sequence[Message]) -> Se
             session_id=self._session.session_id,
             service_session_id=None,
             input_messages=list(input_messages),
+            metadata=self._run_context.metadata,
         )
         for source_id, source_messages in context_messages.items():
             service_call_context.extend_messages(source_id, source_messages)
@@ -1644,6 +1667,15 @@ async def _prepare_service_call_context(self, messages: Sequence[Message]) -> Se
                 context=service_call_context,
                 state=self._session.state.setdefault(provider.source_id, {}),
             )
+        for provider in self._context_providers:
+            if not provider.requires_loaded_context_before_invocation:
+                continue
+            await provider.after_context_loaded(
+                agent=self._agent,
+                session=self._session,
+                context=service_call_context,
+                state=self._session.state.setdefault(provider.source_id, {}),
+            )
         return service_call_context
 
     async def _persist_service_call_response(