Skip to content

Commit 907f7a3

Browse files
Python: add MCP Host history conversion for AG-UI
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 0a71938 commit 907f7a3

9 files changed

Lines changed: 410 additions & 44 deletions

File tree

‎python/packages/ag-ui/AGENTS.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard.
88
- **`AgentFrameworkWorkflow`** - Wraps native `Workflow` objects, or accepts `workflow_factory(thread_id)` for thread-scoped workflow instances without subclassing
99
- **`AGUIChatClient`** - Chat client that speaks AG-UI protocol
1010
- **`AGUIHttpService`** - HTTP service for AG-UI endpoints
11+
- **`agent_framework_messages_to_agui_host_history()`** - Converts persisted Agent Framework messages to bounded
12+
AG-UI Host history while retaining MCP widget payloads and model replay metadata
1113
- **`AGUIEventConverter`** - Converts between Agent Framework and AG-UI events
1214
- **`add_agent_framework_fastapi_endpoint()`** - Add AG-UI endpoint to FastAPI app (`SupportsAgentRun` or `Workflow`)
1315
- **`InMemoryAGUIThreadSnapshotStore`** - Memory-only latest AG-UI Thread Snapshot store for local development, demos, and tests

‎python/packages/ag-ui/agent_framework_ag_ui/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from ._endpoint import add_agent_framework_fastapi_endpoint
1111
from ._event_converters import AGUIEventConverter
1212
from ._http_service import AGUIHttpService
13+
from ._message_adapters import agent_framework_messages_to_agui_host_history
1314
from ._snapshots import (
1415
DEFAULT_MAX_THREAD_SNAPSHOTS,
1516
AGUIThreadID,
@@ -39,6 +40,7 @@
3940
"AgentFrameworkWorkflow",
4041
"WorkflowFactory",
4142
"add_agent_framework_fastapi_endpoint",
43+
"agent_framework_messages_to_agui_host_history",
4244
"AGUIChatClient",
4345
"AGUIChatOptions",
4446
"AGUIEventConverter",

‎python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -102,12 +102,10 @@
102102
)
103103
from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts
104104
from ._utils import (
105-
_AGUI_MCP_TOOL_RESULT_KEY,
106-
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY,
107-
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
108105
_approval_interrupt_id,
109106
_bound_host_payload_history,
110107
_function_call_server_label,
108+
_mcp_host_history_fields,
111109
_model_items_for_agui_replay,
112110
_persistable_host_payload_history,
113111
_project_host_payload_history,
@@ -700,11 +698,10 @@ def _make_approval_tool_result_events(resolved_approval_results: list[Content])
700698
ui_str = _resolve_ui_payload(llm_str, host_payload if has_host_payload else display_result)
701699
replay_properties: dict[str, Any] = {}
702700
if has_host_payload:
703-
replay_properties = {
704-
_AGUI_MCP_TOOL_RESULT_KEY: True,
705-
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: _stringify_tool_result(host_payload),
706-
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: _model_items_for_agui_replay(resolved, llm_str),
707-
}
701+
replay_properties = _mcp_host_history_fields(
702+
host_payload,
703+
_model_items_for_agui_replay(resolved, llm_str),
704+
)
708705
events.append(
709706
ToolCallResultEvent(
710707
message_id=generate_event_id(),
@@ -1996,10 +1993,11 @@ def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) ->
19961993
"content": llm_result,
19971994
}
19981995
if has_host_payload:
1999-
snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True
2000-
snapshot_message[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] = _stringify_tool_result(host_payload)
2001-
snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = _model_items_for_agui_replay(
2002-
content, llm_result
1996+
snapshot_message.update(
1997+
_mcp_host_history_fields(
1998+
host_payload,
1999+
_model_items_for_agui_replay(content, llm_result),
2000+
)
20032001
)
20042002
result_by_call_id[call_id] = snapshot_message
20052003
return result_by_call_id

‎python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py‎

Lines changed: 124 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,26 @@
1717
from agent_framework._types import ContentType # pyright: ignore[reportPrivateUsage]
1818
from typing_extensions import get_args
1919

20+
from ._state import TOOL_RESULT_DISPLAY_KEY
2021
from ._utils import (
2122
_AGUI_HOST_PAYLOAD_OMITTED_KEY,
2223
_AGUI_MCP_TOOL_RESULT_KEY,
2324
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY,
2425
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
26+
_MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES,
2527
AGUI_TO_FRAMEWORK_ROLE,
2628
FRAMEWORK_TO_AGUI_ROLE,
29+
_bound_host_payload_history,
30+
_extract_mcp_tool_result_host_payload,
31+
_extract_tool_result_marker_values,
32+
_host_payload_history_size,
33+
_mcp_host_history_fields,
2734
_model_content_from_mcp_host_payload,
35+
_model_items_for_agui_replay,
36+
_persistable_host_payload_history,
37+
_project_host_payload_history,
2838
_sanitize_model_replay_item,
39+
_stringify_tool_result,
2940
get_role_value,
3041
normalize_agui_role,
3142
safe_json_parse,
@@ -1034,7 +1045,12 @@ def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, A
10341045
return text, tool_calls
10351046

10361047

1037-
def _split_mixed_message_to_agui(msg: Message, role: str, unresolved_call_ids: set[str]) -> list[dict[str, Any]]:
1048+
def _split_mixed_message_to_agui(
1049+
msg: Message,
1050+
role: str,
1051+
unresolved_call_ids: set[str],
1052+
emitted_results: list[tuple[Content, dict[str, Any]]] | None = None,
1053+
) -> list[dict[str, Any]]:
10381054
"""Convert a Message that carries function_result content into ordered AG-UI messages.
10391055
10401056
A single Agent Framework message can interleave assistant content (text,
@@ -1107,14 +1123,15 @@ def flush_segment() -> None:
11071123
messages.append(assistant_msg)
11081124

11091125
def emit_result(content: Content) -> None:
1110-
messages.append(
1111-
{
1112-
"id": next_id(),
1113-
"role": "tool",
1114-
"content": content.result if content.result is not None else "",
1115-
"toolCallId": content.call_id,
1116-
}
1117-
)
1126+
tool_message: dict[str, Any] = {
1127+
"id": next_id(),
1128+
"role": "tool",
1129+
"content": content.result if content.result is not None else "",
1130+
"toolCallId": content.call_id,
1131+
}
1132+
messages.append(tool_message)
1133+
if emitted_results is not None:
1134+
emitted_results.append((content, tool_message))
11181135
if content.call_id is not None:
11191136
unresolved_call_ids.discard(str(content.call_id))
11201137

@@ -1163,15 +1180,13 @@ def drain_queued() -> None:
11631180
return messages
11641181

11651182

1166-
def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]:
1167-
"""Convert Agent Framework messages to AG-UI format.
1168-
1169-
Args:
1170-
messages: List of Agent Framework Message objects or AG-UI dicts (already converted)
1171-
1172-
Returns:
1173-
List of AG-UI message dictionaries
1174-
"""
1183+
def _convert_agent_framework_messages_to_agui(
1184+
messages: list[Message] | list[dict[str, Any]],
1185+
*,
1186+
emitted_results: list[tuple[Content, dict[str, Any]]] | None = None,
1187+
preserve_host_history_dicts: bool = False,
1188+
) -> list[dict[str, Any]]:
1189+
"""Convert Agent Framework messages to AG-UI format."""
11751190
from ._utils import generate_event_id
11761191

11771192
result: list[dict[str, Any]] = []
@@ -1199,6 +1214,12 @@ def track_emitted(
11991214
if isinstance(msg, dict):
12001215
# Always work on a copy to avoid mutating input
12011216
normalized_msg = msg.copy()
1217+
if not preserve_host_history_dicts and normalized_msg.get(_AGUI_MCP_TOOL_RESULT_KEY) is True:
1218+
normalized_msg = _persistable_host_payload_history([normalized_msg])[0].copy()
1219+
normalized_msg.pop(_AGUI_MCP_TOOL_RESULT_KEY, None)
1220+
normalized_msg.pop(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, None)
1221+
normalized_msg.pop(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, None)
1222+
normalized_msg.pop(_AGUI_HOST_PAYLOAD_OMITTED_KEY, None)
12021223
normalized_msg["role"] = normalize_agui_role(normalized_msg.get("role"))
12031224
# Ensure ID exists
12041225
if "id" not in normalized_msg:
@@ -1230,7 +1251,14 @@ def track_emitted(
12301251
# result is dropped and each result stays after its matching call. Messages
12311252
# with no result use the simple single-message form below.
12321253
if any(content.type == "function_result" for content in msg.contents):
1233-
result.extend(_split_mixed_message_to_agui(msg, role, unresolved_call_ids))
1254+
result.extend(
1255+
_split_mixed_message_to_agui(
1256+
msg,
1257+
role,
1258+
unresolved_call_ids,
1259+
emitted_results,
1260+
)
1261+
)
12341262
continue
12351263

12361264
content_text, tool_calls = _encode_agui_segment(msg.contents)
@@ -1250,6 +1278,83 @@ def track_emitted(
12501278
return result
12511279

12521280

1281+
def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]:
1282+
"""Convert Agent Framework messages to model-safe AG-UI request format."""
1283+
return _convert_agent_framework_messages_to_agui(messages)
1284+
1285+
1286+
def _prepare_host_history_fields(
1287+
emitted_results: list[tuple[Content, dict[str, Any]]],
1288+
*,
1289+
max_size_bytes: int,
1290+
) -> tuple[dict[int, dict[str, Any]], set[int]]:
1291+
"""Materialize only the newest MCP Host projections that fit the aggregate budget."""
1292+
retained_fields: dict[int, dict[str, Any]] = {}
1293+
omitted_ids: set[int] = set()
1294+
retained_size = 0
1295+
budget_exhausted = False
1296+
1297+
for content, _ in reversed(emitted_results):
1298+
has_host_payload, host_payload = _extract_mcp_tool_result_host_payload(content)
1299+
if not has_host_payload:
1300+
continue
1301+
content_id = id(content)
1302+
if budget_exhausted:
1303+
omitted_ids.add(content_id)
1304+
continue
1305+
1306+
display_values = _extract_tool_result_marker_values(content, TOOL_RESULT_DISPLAY_KEY)
1307+
if display_values:
1308+
host_payload = display_values[-1]
1309+
model_result = _stringify_tool_result(content.result if content.result is not None else "")
1310+
fields = _mcp_host_history_fields(
1311+
host_payload,
1312+
_model_items_for_agui_replay(content, model_result),
1313+
)
1314+
message_size = _host_payload_history_size({"content": model_result, **fields})
1315+
if retained_size + message_size > max_size_bytes:
1316+
omitted_ids.add(content_id)
1317+
budget_exhausted = True
1318+
continue
1319+
retained_size += message_size
1320+
retained_fields[content_id] = fields
1321+
1322+
return retained_fields, omitted_ids
1323+
1324+
1325+
def agent_framework_messages_to_agui_host_history(
1326+
messages: list[Message] | list[dict[str, Any]],
1327+
*,
1328+
max_host_payload_history_size_bytes: int = _MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES,
1329+
) -> list[dict[str, Any]]:
1330+
"""Convert Agent Framework messages to bounded AG-UI Host history with replay metadata."""
1331+
if messages and isinstance(messages[0], dict):
1332+
converted = _persistable_host_payload_history(
1333+
_convert_agent_framework_messages_to_agui(messages, preserve_host_history_dicts=True)
1334+
)
1335+
else:
1336+
message_objects = cast(list[Message], messages)
1337+
emitted_results: list[tuple[Content, dict[str, Any]]] = []
1338+
converted = _convert_agent_framework_messages_to_agui(
1339+
message_objects,
1340+
emitted_results=emitted_results,
1341+
)
1342+
host_history_fields, omitted_ids = _prepare_host_history_fields(
1343+
emitted_results,
1344+
max_size_bytes=max_host_payload_history_size_bytes,
1345+
)
1346+
for content, tool_message in emitted_results:
1347+
if fields := host_history_fields.get(id(content)):
1348+
tool_message.update(fields)
1349+
elif id(content) in omitted_ids:
1350+
tool_message[_AGUI_HOST_PAYLOAD_OMITTED_KEY] = True
1351+
bounded = _bound_host_payload_history(
1352+
converted,
1353+
max_size_bytes=max_host_payload_history_size_bytes,
1354+
)
1355+
return _project_host_payload_history(bounded)
1356+
1357+
12531358
def extract_text_from_contents(contents: list[Any]) -> str:
12541359
"""Extract text from Agent Framework contents.
12551360

‎python/packages/ag-ui/agent_framework_ag_ui/_run_common.py‎

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,10 @@
3838
from ._predictive_state import PredictiveStateHandler
3939
from ._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY
4040
from ._utils import (
41-
_AGUI_MCP_TOOL_RESULT_KEY,
42-
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY,
43-
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
4441
_approval_interrupt_id,
4542
_extract_mcp_tool_result_host_payload,
4643
_extract_tool_result_marker_values,
44+
_mcp_host_history_fields,
4745
_model_items_for_agui_replay,
4846
_stringify_tool_result,
4947
generate_event_id,
@@ -810,16 +808,11 @@ def _emit_tool_result_common(
810808
}
811809
event_replay_properties: dict[str, Any] = {}
812810
if snapshot_result is not _UNSET:
813-
snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True
814-
snapshot_message[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] = snapshot_result_content
815-
snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = model_items or [
816-
{"type": "text", "text": result_content}
817-
]
818-
event_replay_properties = {
819-
_AGUI_MCP_TOOL_RESULT_KEY: True,
820-
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: snapshot_result_content,
821-
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY],
822-
}
811+
event_replay_properties = _mcp_host_history_fields(
812+
snapshot_result_content,
813+
model_items or [{"type": "text", "text": result_content}],
814+
)
815+
snapshot_message.update(event_replay_properties)
823816
events[-1] = ToolCallResultEvent(
824817
message_id=message_id,
825818
tool_call_id=call_id,

‎python/packages/ag-ui/agent_framework_ag_ui/_utils.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,15 @@ def _host_payload_history_size(message: dict[str, Any]) -> int:
193193
return content_size + sidecar_size
194194

195195

196+
def _mcp_host_history_fields(host_payload: Any, model_items: list[dict[str, Any]]) -> dict[str, Any]:
197+
"""Build the private fields that preserve one MCP Host result for safe replay."""
198+
return {
199+
_AGUI_MCP_TOOL_RESULT_KEY: True,
200+
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: _stringify_tool_result(host_payload),
201+
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: model_items,
202+
}
203+
204+
196205
def _persistable_host_payload_history(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
197206
"""Keep canonical persisted content safe for readers that ignore private replay fields."""
198207
persisted: list[dict[str, Any]] = []

0 commit comments

Comments
 (0)