Python: Preserve MCP Host payloads in AG-UI snapshots - #8129
Python: Preserve MCP Host payloads in AG-UI snapshots#8129Eduard van Valkenburg (eavanvalkenburg) wants to merge 1 commit into
Conversation
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
GitHub Actions
Active
GitHub Actions
Active
GitHub Actions
Active
agent-framework-automation
Bot
added
the
python
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
| "id": message_id, | ||
| "role": "tool", | ||
| "toolCallId": call_id, | ||
| "content": snapshot_result_content, |
There was a problem hiding this comment.
This persists the Host projection in canonical content and depends on two new private fields to recover the model-facing result. During a rolling deployment, an older worker reading a snapshot written by this version ignores those fields and treats the Host JSON as model input; if an explicit display payload contains accepted, the old adapter can even reinterpret it as an approval and re-execute the matching tool. Keep the canonical persisted representation safe for older readers, or version/gate the snapshot format so incompatible workers cannot consume it.
| isinstance(serialized_items, list) | ||
| and serialized_items | ||
| and all( | ||
| isinstance(item, dict) and item.get("type") in {"text", "data", "uri", "error"} |
There was a problem hiding this comment.
The sidecar writer serializes every Content returned by the public MCPTool.parse_tool_results contract, but this allowlist rejects every valid type outside these four. A custom parser returning another model-facing Content type works on the live turn, then the next snapshot replay replaces the entire result with the text fallback and changes provider history. Ensure the writer and reader support the same valid content set, or reject/normalize unsupported items before the first model turn so replay remains equivalent.
| model_text = _model_text_from_replay_items(message) | ||
| if not model_text: | ||
| model_text = _model_content_from_mcp_host_payload(safe_json_parse(message.get("content"))) | ||
| bounded_message["content"] = model_text |
There was a problem hiding this comment.
What happens when the model replay consists only of data, uri, or error items? _model_text_from_replay_items returns an empty string, so eviction substitutes text from the complete Host payload; a custom parser that deliberately excluded sensitive MCP text will then send that Host-only text to the provider on replay. Could the omitted form retain a safe serialization of the actual model items, or use a generic result instead of consulting the Host payload?
| bounded_message["content"] = model_text | ||
| bounded_message.pop(_AGUI_MCP_TOOL_RESULT_KEY, None) | ||
| bounded_message.pop(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, None) | ||
| bounded_message[_AGUI_HOST_PAYLOAD_OMITTED_KEY] = True |
There was a problem hiding this comment.
Should the bounded form retain provenance for model text such as {"accepted": true}? Removing _agentFrameworkMcpResult makes agui_messages_to_agent_framework reinterpret that legitimate terminal result as a new approval response, so it drops the function result and leaves an orphan tool call that provider APIs reject. _agentFrameworkHostPayloadOmitted is never read on the inbound path, so it does not prevent the reclassification.
| serialized_items: list[dict[str, Any]] = [] | ||
| for item in items: | ||
| serialized = item.to_dict() |
There was a problem hiding this comment.
Could replay serialization fail safely before the result is emitted? item.to_dict() recursively traverses arbitrary additional_properties, so a cyclic provider value raises RecursionError after the MCP tool has already run but before AG-UI emits or persists its result; retrying the turn can execute the side effect again. A serialization fallback to the already-bounded model_result would preserve the terminal result without traversing that metadata.
There was a problem hiding this comment.
Could we preserve the model replay metadata on the event path used by AgentFrameworkWorkflow? _emit_tool_result_common puts the complete Host payload in ToolCallResultEvent.content, but adds _agentFrameworkMcpResult and _agentFrameworkModelContent only to flow.tool_results; _WorkflowSnapshotBuilder therefore persists unmarked Host JSON, and the next turn sends structuredContent to the provider. That workflow save path also bypasses _bound_host_payload_history, so these snapshots can grow without the fixed aggregate limit.
2 participants
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
Contentitems 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 does not change the model-content selection behavior established by #7897.
Description & Review Guide
TOOL_CALL_RESULT, approval results, hosted-MCP compatibility results, andMESSAGES_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.TOOL_RESULT_DISPLAY_KEYpayloads remain authoritative, ordinary tool results and MCP-shaped JSON are unchanged, server_metais not replicated into replay items, and existing parallel/mixed function-result conversion remains intact.Related Issue
Part of #7959
This draft is stacked on #8128 and replaces only the AG-UI slice of the existing #7971 implementation; #8128 owns the core payload marker, while this PR owns AG-UI projection and replay safety.
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.