Python: Preserve FIDES AG-UI replay provenance

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

@@ -8,7 +8,7 @@
 import json
 import logging
 import uuid
-from collections.abc import AsyncIterable, Awaitable, Collection, Mapping, Sequence
+from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Mapping, Sequence
 from dataclasses import dataclass, field
 from functools import partial
 from typing import TYPE_CHECKING, Any, cast
@@ -61,7 +61,7 @@
     _try_execute_function_call_groups,  # type: ignore
     normalize_function_invocation_configuration,
 )
-from agent_framework._types import ResponseStream
+from agent_framework._types import APPROVAL_REJECTION_RESULT_KEY, ResponseStream
 from agent_framework.exceptions import AgentInvalidResponseException
 from agent_framework.observability import (
     _use_telemetry_conversation_id,  # pyright: ignore[reportPrivateUsage]
@@ -90,7 +90,13 @@
     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,
+    _copy_backend_snapshot_origin,
+    _mark_backend_snapshot_origin,
+    _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 (
@@ -1764,6 +1770,7 @@ async def _resolve_approval_responses(
     pending_response_groups: dict[object, tuple[ApprovalOccurrence, list[Content]]] = {}
     intents_by_response_content_id: dict[int, AuthorizedExecution] = {}
     authenticated_non_grants: list[Content] = []
+    authenticated_rejection_occurrence_ids: set[str] = set()
     for response in approval_responses:
         resp_id = response.id
         function_call_id = response.function_call.call_id if response.function_call else None
@@ -1861,6 +1868,7 @@ async def _resolve_approval_responses(
                 )
         if decision_is_boolean and primary_response.approved is False and not server_label:
             authenticated_non_grants.append(_approval_observer_response(pending_entry))
+            authenticated_rejection_occurrence_ids.add(pending_entry.function_call_id)
         if (
             (primary_response.approved is True or pending_entry.owner is ApprovalExecutionOwner.DEFERRED)
             and lifecycle is not None
@@ -1902,6 +1910,14 @@ async def _resolve_approval_responses(
             authenticated_non_grants, session=invocation_session
         )
 
+    if valid_response_content_ids:
+        for message in messages:
+            if message.contents and all(
+                content.type == "function_approval_response" and id(content) in valid_response_content_ids
+                for content in message.contents
+            ):
+                _mark_backend_snapshot_origin(message)
+
     if response_content_ids_to_strip:
         filtered_messages: list[Message] = []
         for message in messages:
@@ -2069,18 +2085,32 @@ async def execute_local_call(approval: Content = approval, call_id: str = call_i
             if id(response) in deferred_response_content_ids
         }
     )
-    _replace_approval_contents_with_results(messages, normalization_responses, replacement_groups)
+    resolved_contents = _replace_approval_contents_with_results(messages, normalization_responses, replacement_groups)
+    authenticated_rejection_result_ids = {
+        id(content)
+        for content in resolved_contents
+        if content.type == "function_result"
+        and content.id in authenticated_rejection_occurrence_ids
+        and content.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is True
+    }
 
     # Post-process: Convert user messages with function_result content to proper tool messages.
     # After _replace_approval_contents_with_results, approved tool calls have their results
     # placed in user messages. OpenAI requires tool results to be in role="tool" messages.
     # This transformation ensures the message history is valid for the LLM provider.
-    _convert_approval_results_to_tool_messages(messages)
+    _convert_approval_results_to_tool_messages(
+        messages,
+        backend_result_content_ids=authenticated_rejection_result_ids,
+    )
 
     return approved_results
 
 
-def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None:
+def _convert_approval_results_to_tool_messages(
+    messages: list[Message],
+    *,
+    backend_result_content_ids: set[int] | None = None,
+) -> None:
     """Convert function_result content in user messages to proper tool messages.
 
     After approval processing, tool results end up in user messages. OpenAI and other
@@ -2091,11 +2121,25 @@ def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None:
 
     Args:
         messages: List of Message objects to process
+        backend_result_content_ids: Object identities of exact backend-synthesized
+            results whose derived tool messages retain snapshot provenance.
     """
     result: list[Message] = []
+    backend_result_ids = backend_result_content_ids or set()
 
     for msg in messages:
         if get_role_value(msg) != "user":
+            backend_contents = [content for content in msg.contents if id(content) in backend_result_ids]
+            if backend_contents and len(backend_contents) == len(msg.contents):
+                _mark_backend_snapshot_origin(msg)
+            elif backend_contents and get_role_value(msg) == "tool":
+                for content in msg.contents:
+                    split_message = Message(role="tool", contents=[content])
+                    _copy_backend_snapshot_origin(msg, split_message)
+                    if id(content) in backend_result_ids:
+                        _mark_backend_snapshot_origin(split_message)
+                    result.append(split_message)
+                continue
             result.append(msg)
             continue
 
@@ -2113,11 +2157,17 @@ def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None:
 
         # Tool messages first (right after the preceding assistant message per OpenAI requirements)
         for func_result in function_results:
-            result.append(Message(role="tool", contents=[func_result]))
+            tool_message = Message(role="tool", contents=[func_result])
+            _copy_backend_snapshot_origin(msg, tool_message)
+            if id(func_result) in backend_result_ids:
+                _mark_backend_snapshot_origin(tool_message)
+            result.append(tool_message)
 
         # Then user message with remaining content (if any)
         if other_contents:
-            result.append(Message(role="user", contents=other_contents))
+            user_message = Message(role="user", contents=other_contents)
+            _copy_backend_snapshot_origin(msg, user_message)
+            result.append(user_message)
 
     messages[:] = result
 
@@ -2801,6 +2851,47 @@ 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 _collect_fides_replay_folders(agent: SupportsAgentRun) -> list[Callable[..., None]]:
+    """Return replay-folding hooks exposed by FIDES context providers."""
+    folders: list[Callable[..., None]] = []
+    for provider in getattr(agent, "context_providers", ()):
+        fold = getattr(provider, "_fold_uncovered_replay_messages", None)
+        if callable(fold):
+            folders.append(cast(Callable[..., None], fold))
+    return folders
+
+
+def _fold_fides_uncovered_replay_messages(
+    folders: Sequence[Callable[..., None]],
+    session: AgentSession,
+    messages: Sequence[Message],
+    *,
+    complete_replay_messages: Sequence[Message] | None = None,
+    replay_history_complete: bool,
+) -> None:
+    """Fold uncovered replay into FIDES state before a snapshot-only return."""
+    for fold in folders:
+        fold(
+            session,
+            messages,
+            complete_replay_messages=complete_replay_messages,
+            replay_history_complete=replay_history_complete,
+        )
+
+
 def _legacy_tool_message_approval_resume(
     messages: Sequence[Mapping[str, Any]],
     *,
@@ -3004,6 +3095,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 +3165,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 +3292,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 +3457,29 @@ async def save_thread_snapshot(
             ),
         )
 
+    fides_replay_folders = _collect_fides_replay_folders(agent)
+
+    def fold_fides_replay_before_snapshot_save() -> None:
+        if not fides_replay_folders:
+            return
+        complete_replay_messages: Sequence[Message] | None = None
+        if not config.use_service_session:
+            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,
+                protected_tool_call_ids=protected_tool_call_ids,
+            )
+        _fold_fides_uncovered_replay_messages(
+            fides_replay_folders,
+            session,
+            messages,
+            complete_replay_messages=complete_replay_messages,
+            replay_history_complete=not config.use_service_session,
+        )
+        _mark_fides_state_covered_messages(agent, messages)
+
     authenticated_cancellations = [
         _approval_observer_response(occurrence, cancelled=True)
         for interrupt_id in cancelled_resume_ids
@@ -3552,6 +3671,7 @@ async def save_thread_snapshot(
         forwarded_executions=forwarded_executions,
     )
     resolved_approval_results = retained_approval_results + newly_resolved_approval_results
+    _mark_fides_state_covered_messages(agent, _take_backend_snapshot_messages(messages))
 
     # Defense-in-depth: replace approval payloads in snapshot with actual tool results
     # so CopilotKit does not re-send stale approval content on subsequent turns.
@@ -3613,6 +3733,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 +3767,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 +3776,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(