Skip to content

Python: Preserve MCP Host payloads in AG-UI snapshots #8129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Eduard van Valkenburg (eavanvalkenburg) wants to merge 1 commit into from
Draft

Python: Preserve MCP Host payloads in AG-UI snapshots #8129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 24 additions & 8 deletions python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,8 +91,8 @@
_new_tool_call_segment_id, # type: ignore
_reconstruct_messages_from_thread_snapshot, # type: ignore
_resume_contract_error, # type: ignore
_resolve_tool_result_host_payload, # type: ignore
_resolve_ui_payload, # type: ignore
_stringify_tool_result, # type: ignore
_track_tool_call_segment, # type: ignore
)
from ._snapshots import (
Expand All @@ -102,8 +102,13 @@
)
from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts
from ._utils import (
_AGUI_MCP_TOOL_RESULT_KEY,
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
_approval_interrupt_id,
_bound_host_payload_history,
_function_call_server_label,
_model_items_for_agui_replay,
_stringify_tool_result,
canonical_function_arguments,
convert_agui_tools_to_agent_framework,
generate_event_id,
Expand Down Expand Up @@ -687,7 +692,9 @@ def _make_approval_tool_result_events(resolved_approval_results: list[Content])
if resolved.call_id:
raw = resolved.result if resolved.result is not None else ""
llm_str = _stringify_tool_result(raw)
ui_str = _resolve_ui_payload(llm_str, _extract_tool_result_display(resolved))
display_result = _extract_tool_result_display(resolved)
has_host_payload, host_payload = _resolve_tool_result_host_payload(resolved, display_result)
ui_str = _resolve_ui_payload(llm_str, host_payload if has_host_payload else display_result)
events.append(
ToolCallResultEvent(
message_id=generate_event_id(),
Expand Down Expand Up @@ -1968,12 +1975,21 @@ def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) ->
]
for content in function_results:
call_id = str(content.call_id)
result_by_call_id[call_id] = {
llm_result = _stringify_tool_result(content.result if content.result is not None else "")
display_result = _extract_tool_result_display(content)
has_host_payload, host_payload = _resolve_tool_result_host_payload(content, display_result)
snapshot_message: dict[str, Any] = {
"id": msg.message_id if msg.message_id and len(function_results) == 1 else generate_event_id(),
"role": "tool",
"toolCallId": call_id,
"content": _stringify_tool_result(content.result if content.result is not None else ""),
"content": _stringify_tool_result(host_payload) if has_host_payload else llm_result,
}
if has_host_payload:
snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True
snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = _model_items_for_agui_replay(
content, llm_result
)
result_by_call_id[call_id] = snapshot_message
return result_by_call_id


Expand Down Expand Up @@ -2067,7 +2083,7 @@ def _build_messages_snapshot(

if flow.snapshot_segments:
_append_segmented_snapshot_messages(flow, all_messages)
return MessagesSnapshotEvent(messages=all_messages) # type: ignore[arg-type]
return MessagesSnapshotEvent(messages=_bound_host_payload_history(all_messages)) # type: ignore[arg-type]

# Add assistant message with tool calls only (no content)
if flow.pending_tool_calls:
Expand Down Expand Up @@ -2101,7 +2117,7 @@ def _build_messages_snapshot(
# MESSAGES_SNAPSHOT retain reasoning content after streaming ends.
all_messages.extend(flow.reasoning_messages)

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


def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str, Any]]:
Expand Down Expand Up @@ -2876,7 +2892,7 @@ async def run_agent_stream(
# stored history unless this run already seeded raw messages from it.
persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages)
await snapshot_session.save(
messages=persisted_messages,
messages=_bound_host_payload_history(persisted_messages),
state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None,
interrupt=None,
session_state=_safe_serialize_session_continuation_state(
Expand Down Expand Up @@ -3237,7 +3253,7 @@ async def run_agent_stream(
# stored history unless this run already seeded raw messages from it.
persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages)
await snapshot_session.save(
messages=persisted_messages,
messages=_bound_host_payload_history(persisted_messages),
state=latest_state_snapshot,
interrupt=flow.interrupts or None,
session_state=_safe_serialize_session_continuation_state(
Expand Down
29 changes: 29 additions & 0 deletions python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,11 @@
)

from ._utils import (
_AGUI_MCP_TOOL_RESULT_KEY,
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
AGUI_TO_FRAMEWORK_ROLE,
FRAMEWORK_TO_AGUI_ROLE,
_model_content_from_mcp_host_payload,
get_role_value,
normalize_agui_role,
safe_json_parse,
Expand Down Expand Up @@ -719,6 +722,32 @@ def _filter_modified_args(
elif isinstance(result_content, dict):
parsed = cast(dict[str, Any], result_content)

if msg.get(_AGUI_MCP_TOOL_RESULT_KEY) is True:
serialized_items = msg.get(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY)
model_items: list[Content] | None = None
if (
isinstance(serialized_items, list)
and serialized_items
and all(
isinstance(item, dict) and item.get("type") in {"text", "data", "uri", "error"}

@github-actions github-actions Bot Sep 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

for item in serialized_items
)
):
try:
model_items = [Content.from_dict(item) for item in serialized_items]
except (TypeError, ValueError):
model_items = None
if not model_items:
model_items = [Content.from_text(_model_content_from_mcp_host_payload(parsed))]
chat_msg = Message(
role="tool",
contents=[Content.from_function_result(call_id=str(tool_call_id), result=model_items)],
)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
continue

is_approval = parsed is not None and "accepted" in parsed

if is_approval:
Expand Down
78 changes: 49 additions & 29 deletions python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
View file
Open in desktop

@moonbox3 Evan Mattson (moonbox3) Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,18 @@

from ._predictive_state import PredictiveStateHandler
from ._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY
from ._utils import _approval_interrupt_id, generate_event_id, make_json_safe, normalize_agui_role
from ._utils import (
_AGUI_MCP_TOOL_RESULT_KEY,
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
_approval_interrupt_id,
_extract_mcp_tool_result_host_payload,
_extract_tool_result_marker_values,
_model_items_for_agui_replay,
_stringify_tool_result,
generate_event_id,
make_json_safe,
normalize_agui_role,
)

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -696,22 +707,6 @@ def _emit_tool_call(
return events


def _extract_tool_result_marker_values(content: Content, key: str) -> list[Any]:
"""Extract marker values from outer and inner tool-result content."""
values: list[Any] = []

outer_ap = getattr(content, "additional_properties", None) or {}
if key in outer_ap:
values.append(outer_ap[key])

for item in content.items or ():
item_ap = getattr(item, "additional_properties", None) or {}
if key in item_ap:
values.append(item_ap[key])

return values


def _extract_tool_result_state(content: Content) -> dict[str, Any] | None:
"""Extract a deterministic AG-UI state update from a tool-result ``Content``.

Expand Down Expand Up @@ -745,15 +740,19 @@ def _extract_tool_result_display(content: Content) -> Any: # noqa: ANN401
return display_values[-1] if display_values else _UNSET


def _stringify_tool_result(raw_result: Any) -> str: # noqa: ANN401
return raw_result if isinstance(raw_result, str) else json.dumps(make_json_safe(raw_result))


def _resolve_ui_payload(llm_str: str, display_result: Any) -> str: # noqa: ANN401
"""Pick the UI-bound string: the serialized display payload when set, else the LLM string."""
return llm_str if display_result is _UNSET else _stringify_tool_result(display_result)


def _resolve_tool_result_host_payload(content: Content, display_result: Any) -> tuple[bool, Any]: # noqa: ANN401
"""Resolve an MCP Host payload, preferring an explicit display projection."""
has_host_payload, host_payload = _extract_mcp_tool_result_host_payload(content)
if has_host_payload and display_result is not _UNSET:
host_payload = display_result
return has_host_payload, host_payload


def _emit_tool_result_common(
call_id: str,
raw_result: Any,
Expand All @@ -762,6 +761,8 @@ def _emit_tool_result_common(
*,
state_update: Mapping[str, Any] | None = None,
display_result: Any = _UNSET, # noqa: ANN401
snapshot_result: Any = _UNSET, # noqa: ANN401
model_items: list[dict[str, Any]] | None = None,
) -> list[BaseEvent]:
"""Shared helper for emitting ToolCallEnd + ToolCallResult events and performing FlowState cleanup.

Expand Down Expand Up @@ -789,6 +790,7 @@ def _emit_tool_result_common(

result_content = _stringify_tool_result(raw_result)
ui_result_content = _resolve_ui_payload(result_content, display_result)
snapshot_result_content = _resolve_ui_payload(result_content, snapshot_result)
message_id = generate_event_id()
events.append(
ToolCallResultEvent(
Expand All @@ -799,14 +801,18 @@ def _emit_tool_result_common(
)
)

flow.tool_results.append(
{
"id": message_id,
"role": "tool",
"toolCallId": call_id,
"content": result_content,
}
)
snapshot_message: dict[str, Any] = {
"id": message_id,
"role": "tool",
"toolCallId": call_id,
"content": snapshot_result_content,

@github-actions github-actions Bot Sep 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

}
if snapshot_result is not _UNSET:
snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True
snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = model_items or [
{"type": "text", "text": result_content}
]
flow.tool_results.append(snapshot_message)
# A result closes the current tool-call segment: a later call opens a new
# one, so `call A -> result A -> call B` snapshots as two call/result pairs
# in stream order instead of grouping B with A (moonbox3's replay concern).
Expand Down Expand Up @@ -851,13 +857,20 @@ def _emit_tool_result(
raw_result = content.result if content.result is not None else ""
state_update = _extract_tool_result_state(content)
display_result = _extract_tool_result_display(content)
has_host_payload, host_payload = _resolve_tool_result_host_payload(content, display_result)
if has_host_payload and display_result is _UNSET:
display_result = host_payload
return _emit_tool_result_common(
content.call_id,
raw_result,
flow,
predictive_handler,
state_update=state_update,
display_result=display_result,
snapshot_result=host_payload if has_host_payload else _UNSET,
model_items=(
_model_items_for_agui_replay(content, _stringify_tool_result(raw_result)) if has_host_payload else None
),
)


Expand Down Expand Up @@ -1022,13 +1035,20 @@ def _emit_mcp_tool_result(
raw_output = content.output if content.output is not None else ""
state_update = _extract_tool_result_state(content)
display_result = _extract_tool_result_display(content)
has_host_payload, host_payload = _resolve_tool_result_host_payload(content, display_result)
if has_host_payload and display_result is _UNSET:
display_result = host_payload
return _emit_tool_result_common(
content.call_id,
raw_output,
flow,
predictive_handler,
state_update=state_update,
display_result=display_result,
snapshot_result=host_payload if has_host_payload else _UNSET,
model_items=(
_model_items_for_agui_replay(content, _stringify_tool_result(raw_output)) if has_host_payload else None
),
)


Expand Down
Loading
Loading