Repository navigation
Python: Preserve FIDES AG-UI replay provenance - #9878
Eduard van Valkenburg (eavanvalkenburg) wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
MAF Automated Review — Iteration 1
Result: Findings reported
Scope: full PR (1 commit(s)): c127b36c74d5
Model: gpt-5.6-sol
Overview
The PR establishes a strong object-identity boundary for backend snapshot provenance, preserves it through normalization, and folds uncovered replay before the new snapshot-only save points. The regression coverage exercises restored history, missing labels, confirmation, authenticated rejection, and provider-failure paths. Two residual trust-boundary issues remain: the supported AG-UI dependency range permits a core version without the new security seam, and authenticated rejection provenance is assigned by a reusable call ID and public result string rather than by the exact lifecycle-produced result.
Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
2 verified findings remained after source verification (1 high, 1 medium) across 1 file. Details are attached to the affected lines below.
Affected areas: python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
| if not messages: | ||
| return | ||
| for provider in getattr(agent, "context_providers", ()): | ||
| marker = getattr(provider, "_mark_replay_state_covered_messages", None) |
There was a problem hiding this comment.
agent-framework-ag-ui still permits core 1.22, but these replay methods are introduced in core 1.23, so this optional lookup silently disables the security fix for a supported installation. On a snapshot-only approval/confirmation return, fresh client replay can then be persisted while the previous trusted FIDES label remains, and approval-time execution can consume that stale label before any later before_run reconstruction. Please either raise the core dependency floor to 1.23 or fail closed when a FIDES provider lacks this seam.
| for msg in messages: | ||
| if authenticated_rejection_call_ids and any( | ||
| content.type == "function_result" | ||
| and content.call_id in authenticated_rejection_call_ids |
There was a problem hiding this comment.
This authenticates any result with the same reusable call ID and public rejection string, not the exact result synthesized from the validated lifecycle response. A request can submit a valid denial together with a client-owned matching tool result; that message is then marked backend-origin and excluded from FIDES replay reconstruction, allowing fresh client replay to retain the restored trusted label. Please carry the exact synthesized result identity/occurrence through normalization and mark only that result.
Code coverage overviewLanguages: Python Python / code-coverage/python
View coverage of the 6 most impacted files.
ℹ️ Calculations represent line coverage. Results may change as calculations are updated. Updated |
c127b36 to
74b7133
Compare
There was a problem hiding this comment.
🟡 Changes recommended
Active security flaws can preserve trusted provenance for client-controlled tool data and reconstruct labels from incomplete history, bypassing conservative labeling and fail-closed protections.
4 open issues
What changed in this PR
This PR adds the AG-UI replay-provenance layer of the FIDES hardening stack, distinguishing backend-restored history from fresh client data when restoring and persisting security state.
Changes:
- Preserve snapshot provenance through message normalization and remove private markers before provider exposure.
- Fold uncovered replay before confirmation and approval snapshot saves, including missing-label reconstruction.
- Add regression coverage for replay provenance, authenticated denials, and provider-failure persistence.
| File | Description |
|---|---|
python/packages/core/tests/test_security.py |
Tests restored-state handling and replay-marker removal. |
python/packages/core/agent_framework/security.py |
Adds transport-aware replay labeling and snapshot-time folding. |
python/packages/ag-ui/tests/ag_ui/test_endpoint.py |
Adds end-to-end provenance and security-state persistence tests. |
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py |
Carries snapshot provenance through message conversion and sanitation. |
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py |
Connects snapshot and approval provenance to FIDES folding before saves. |
Review effort: Balanced
| if authenticated_rejection_call_ids and any( | ||
| content.type == "function_result" | ||
| and content.call_id in authenticated_rejection_call_ids | ||
| and content.result == "Error: Tool call invocation was rejected by user." | ||
| for content in msg.contents | ||
| ): | ||
| _mark_backend_snapshot_origin(msg) |
| if backend_snapshot_message_ids is not None and id(raw_message) in backend_snapshot_message_ids: | ||
| message.additional_properties[_BACKEND_SNAPSHOT_ORIGIN_KEY] = _BACKEND_SNAPSHOT_ORIGIN_MARKER |
| replay_messages = ( | ||
| messages if restored_context_label or complete_replay_messages is None else complete_replay_messages | ||
| ) |
| ) | ||
|
|
||
| def fold_fides_replay_before_snapshot_save() -> None: | ||
| complete_replay_snapshot = snapshot_messages |


Motivation & Context
FIDES security state can be restored from an AG-UI Thread Snapshot while the current turn adds a fresh client tool
result or resolves an approval. Without transport provenance, normalization can either reclassify restored history as
fresh external data or let fresh data inherit coverage from the stored snapshot. Confirmation and approval safe-point
returns can also persist before FIDES folds new replay data.
This is the fourth layer of the upstream replacement stack for #9652. It builds on #9875, #9876, and #9877 and
isolates AG-UI replay provenance, snapshot folding, and authenticated approval-result handling from the later MCP
carrier layer.
Description & Review Guide
preserves that origin through reasoning removal, tool-history sanitation, deduplication, and synthetic skipped
results, and clears the transport marker before provider exposure. A minimal FIDES replay seam treats covered
messages as represented by restored state, folds uncovered messages before snapshot-only or approval safe-point
saves, and reconstructs complete current-format replay when the stored context label is missing. Lifecycle-
authenticated approvals, denials, rejection results, replacement approvals, and mixed siblings retain backend
authority, while fresh client results remain conservative.
fresh client tool results cannot inherit snapshot coverage, and confirmation or provider-failure paths persist the
correct conservative FIDES label. There are no public API or breaking changes. Validation passed with:
cd python && uv run pytest -q packages/ag-ui/tests/ag_ui/test_endpoint.py -k 'fides_snapshot_provenance_survives_reasoning_normalization or fides_synthetic_skipped_result_keeps_snapshot_provenance or fides_normal_continuation_reconstructs_missing_context_label or fides_confirm_changes_folds_fresh_client_result or fides_confirm_changes_reconstructs_missing_context_label or fides_authenticated_rejection_preserves_restored_context or fides_rejection_folds_client_result_before_provider_failure' --tb=short(9 passed, 441 deselected)cd python && uv run pytest -q packages/core/tests/test_security.py -k 'transport_marked_replay_uses_restored_state_and_clears_marker' --tb=short(1 passed)cd python && uv run poe syntax -P ag-ui -C(passed)cd python && uv run poe pyright -P ag-ui(0 errors and 0 warnings)cd python && uv run poe test-typing -P ag-ui(mypy, pyrefly, pyright, ty, and zuban passed)cd python && uv run poe test -P ag-ui(1,786 passed)cd python && uv run poe syntax -P core -C(passed)cd python && uv run poe pyright -P core(0 errors and 0 warnings)cd python && uv run poe test-typing -P core(mypy, pyrefly, pyright, ty, and zuban passed)cd python && uv run poe test -P core(9,163 passed, 20 skipped, 2 xfailed)normalization and deduplication, the ordering of FIDES folding before confirmation and approval safe-point saves,
and authenticated rejection or replacement provenance. This layer intentionally excludes service-history ownership,
general replay or approval binding, policy materialization, and MCP IFC changes.
Related Issue
Fixes #9652
#9652 is the original source draft. This PR is its AG-UI replay-provenance replacement layer, stacked as
#9875 -> #9876 -> #9877 -> this PR rather than a competing implementation.
AI Assistance
AI assistance details: GitHub Copilot assisted with implementation, regression tests, validation, and this PR
description.
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.