Repository navigation
Python: Preserve FIDES AG-UI replay provenance - #9878
Eduard van Valkenburg (eavanvalkenburg) wants to merge 2 commits 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.
This layer depends on the core replay seams introduced after
agent-framework-core1.22. The normal version-bump PRwill update the AG-UI package's core dependency floor when that core release is available; this stacked PR
intentionally does not make a release-floor change.
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, reconstructs complete local replay when the stored context label is missing, and rejects reconstruction when
service-owned history is unavailable. Lifecycle-authenticated approvals, replacement approvals, and mixed siblings
retain backend authority. Denial provenance binds to the exact synthesized result object, edited client calls lose
snapshot provenance, and fresh client results remain conservative. Non-FIDES approval paths avoid the extra full
transcript conversion.
fresh client tool results cannot inherit snapshot coverage, and confirmation or provider-failure paths persist the
correct conservative FIDES label. Missing FIDES state fails closed for service-managed history rather than trusting
an incomplete AG-UI snapshot. 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_service_confirmation_does_not_reconstruct_incomplete_history or fides_rejection_does_not_cover_client_result_with_reused_call_id or fides_rejected_edited_approval_invalidates_snapshot_provenance or confirmation_without_fides_does_not_rebuild_snapshot or fides_authenticated_rejection_preserves_restored_context or 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_rejection_folds_client_result_before_provider_failure' --tb=short(13 passed, 441 deselected)cd python && uv run pytest -q packages/ag-ui/tests/ag_ui/test_message_hygiene.py -k 'convert_approval_results' --tb=short(3 passed, 7 deselected)cd python && uv run pytest -q packages/core/tests/test_security.py -k 'transport_fold_rejects_missing_label_for_incomplete_service_history or transport_marked_replay_uses_restored_state_and_clears_marker' --tb=short(2 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,791 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,276 passed, 20 skipped, 2 xfailed)normalization and deduplication, the ordering of FIDES folding before confirmation and approval safe-point saves,
exact synthesized-denial provenance, edited-call invalidation, and the fail-closed service-history boundary. 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.