Python: Preserve MCP Host payloads in AG-UI snapshots

[Eduard van Valkenburg (eavanvalkenburg)]

Motivation & Context

AG-UI live tool events and complete message snapshots currently expose only the model-facing MCP result, so Host applications lose structured content and other complete MCP result fields. Snapshot replay also needs to preserve the original model-facing Content items without allowing Host/UI JSON to enter provider history.

This is layer 2 of the #7971 replacement stack. It depends on #8128, which introduces the bounded private core marker carrying the complete MCP Host payload. This layer also contains focused follow-up fixes for late #8128 review findings so custom MCP parsers cannot promote server IFC metadata, retained metadata shares the request budget, approval budget state remains serializable, and parser failures still preserve the raw Host carrier. It does not change the model-content selection behavior established by #7897.

Description & Review Guide

• What are the major changes? Detect the layer-1 marker with an older-core fallback; project complete MCP Host payloads through live TOOL_CALL_RESULT, approval results, hosted-MCP compatibility results, and MESSAGES_SNAPSHOT; attach private, lossless model-content replay metadata; restore that metadata only on marked inbound snapshot messages; and enforce a fixed aggregate history budget covering both Host JSON and replay sidecars while retaining newest projections. The focused core follow-up isolates server _meta.ifc from custom parser projections, charges both retained _meta and complete Host payloads to the shared request budget, persists only JSON-compatible approval budget counters, and captures raw MCP results before normal, fallback, or completed-task parsing.
• What is the impact of these changes? AG-UI Hosts receive structured MCP result data consistently across live and replayed flows, while model history remains unchanged and safe. Explicit TOOL_RESULT_DISPLAY_KEY payloads remain authoritative, ordinary tool results and MCP-shaped JSON are unchanged, server _meta is not replicated into replay items or custom parser model projections, built-in MCP parsing retains its existing IFC behavior, and existing parallel/mixed function-result conversion remains intact.
• What do you want reviewers to focus on? The provenance gate for inbound replay, lossless rich-content restoration, approval/live/snapshot parity, newest-first aggregate retention after persisted-history seeding, custom-parser IFC trust boundary, and serializable request-budget continuity across approval resume.

Related Issue

Part of #7959

This draft is stacked on #8128 and replaces the AG-UI slice of the existing #7971 implementation. #8128 owns the core payload marker; this PR owns AG-UI projection and replay safety plus the focused late core review follow-up described above.

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

[github-actions]

MAF Automated Review — Iteration 1

Result: Findings reported
Scope: full PR (1 commit(s)): 9289ea07a267
Model: gpt-5.6-sol-fast

Overview

The change consistently projects complete MCP Host payloads across live, approval, hosted-MCP, and snapshot paths while preserving model-facing content in a private replay sidecar. Exact marker checks, metadata stripping, safe fallback behavior, and newest-first aggregate retention provide strong guardrails. Two compatibility gaps remain: persisted snapshots are unsafe for older readers in a rolling deployment, and valid custom MCP Content types are silently degraded on replay.

Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
2 verified findings remained after source verification (2 medium) across 2 files. Details are attached to the affected lines below.

Affected areas: python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py, python/packages/ag-ui/agent_framework_ag_ui/_run_common.py

[Copilot]

🟡 Changes recommended

Unresolved critical and moderate issues can drop Host payloads, alter replay content, bypass budgets, or break serialization.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Preserves complete MCP Host payloads in AG-UI snapshots while isolating model-facing history.

Changes:

• Adds Host-payload projection and bounded persistence.
• Adds lossless model-content replay metadata.
• Covers live, approval, workflow, hydration, and replay paths.

File summaries

File	Review
python/packages/ag-ui/tests/ag_ui/test_snapshots.py	Tests workflow snapshot replay metadata.
python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py	Tests Host-payload hydration.
python/packages/ag-ui/tests/ag_ui/test_run_common.py	Tests projection, budgeting, and compatibility.
python/packages/ag-ui/tests/ag_ui/test_message_adapters.py	Tests safe inbound replay.
python/packages/ag-ui/agent_framework_ag_ui/_workflow.py	Persists bounded workflow snapshots. No issues found.
python/packages/ag-ui/agent_framework_ag_ui/_utils.py	Moderate (1 vote each): Preserve valid empty content projections; make serialized items JSON-safe; include non-string Host payloads in aggregate budget calculations.
python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py	Projects Host payloads during hydration. No issues found.
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py	Moderate (1 vote): Use an explicit None check so valid empty model projections remain empty.
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py	Moderate (1 vote): Accept empty replay sidecars and fall back only when deserialization returns None.
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py	Critical (1 vote): Approval merging can replace earlier marked MCP snapshots with model-only entries, permanently dropping Host payloads. Preserve existing marked entries when replacements lack Host projections. Nit (1 vote): Update the required function-loop specification and scenario-to-test mapping.

Review details

Suppressed comments (2)

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

• This establishes new normative transport behavior for caller-visible MCP results, persisted history, and model-bound replay, but the required function-loop specification has no corresponding Host projection/replay or budget scenario. python/AGENTS.md:62-67 requires the smallest affected sections of docs/specs/004-python-function-calling-loop.md to be updated when normative behavior changes, including the scenario-to-test mapping.

        bounded_messages = _bound_host_payload_history(_persistable_host_payload_history(all_messages))
        return MessagesSnapshotEvent(messages=_project_host_payload_history(bounded_messages))  # type: ignore[arg-type]

python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py:745

• Checking only the type discriminator does not validate the rest of an untrusted replay item. For example, {"type": "text", "text": 1} passes this gate and Content.from_dict, but Content.from_function_result then raises TypeError while joining the non-string text, aborting the inbound run instead of using the safe Host fallback. Keep deserialization and function-result construction inside the guarded path and fall back for malformed values.

                    and all(
                        isinstance(item, dict) and item.get("type") in _VALID_CONTENT_TYPES for item in serialized_items
                    )
                ):
                    try:
                        model_items = [
                            Content.from_dict(_sanitize_model_replay_item(item)) for item in serialized_items

• Files reviewed: 10/10 changed files
• Comments generated: 6
• Review effort level: Balanced

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

[github-code-quality]

Code Coverage Overview

Languages: Python

Python / code-coverage/python

The overall line coverage in commit 58f5570 in the ag-ui-mcp-live-snaps... branch is 91%. Line coverage data for the main branch is not yet available.

Show a line coverage summary of the most covered files.

File	main	ag-ui-mcp-live-snaps... 58f5570	+/-
packages/core/a...work/_skills.py	—	95%	—
packages/core/a...ework/_tools.py	—	94%	—
packages/core/a.../_compaction.py	—	94%	—
packages/core/a...rk/_sessions.py	—	94%	—
packages/core/a...ework/_types.py	—	93%	—
packages/core/a...bservability.py	—	93%	—
packages/openai..._chat_client.py	—	92%	—
packages/core/a...amework/_mcp.py	—	91%	—
packages/ag-ui/...i/_agent_run.py	—	91%	—
packages/foundr...g/_responses.py	—	87%	—

---

_{Updated September 08, 2026 13:20 UTC}

[github-actions]

MAF Automated Review — Iteration 2

Result: Findings reported
Scope: full PR (1 commit(s)): 14a83a5c1b1c
Model: gpt-5.6-sol-fast

Overview

The change separates complete MCP Host payloads from model-facing replay content, applies aggregate retention limits, and preserves that state across approval and workflow snapshot paths. Provenance checks, canonical model-safe persistence, IFC isolation for normal custom-parser output, and focused parity tests provide strong guardrails. One malformed-input path remains: deeply nested replay sidecars can escape the deserialization fallback and abort an inbound run.

Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
1 verified finding remained after source verification (1 medium) across 1 file. Details are attached to the affected lines below.

Affected areas: python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py