Skip to content

Commit 34185de

Browse files
Python: add MCP Host history conversion for AG-UI
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 2034c64 commit 34185de

9 files changed

Lines changed: 538 additions & 42 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 & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -103,11 +103,10 @@
103103
from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts
104104
from ._utils import (
105105
_AGUI_MCP_TOOL_RESULT_KEY,
106-
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY,
107-
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
108106
_approval_interrupt_id,
109107
_bound_host_payload_history,
110108
_function_call_server_label,
109+
_mcp_host_history_fields,
111110
_model_items_for_agui_replay,
112111
_persistable_host_payload_history,
113112
_project_host_payload_history,
@@ -700,11 +699,10 @@ def _make_approval_tool_result_events(resolved_approval_results: list[Content])
700699
ui_str = _resolve_ui_payload(llm_str, host_payload if has_host_payload else display_result)
701700
replay_properties: dict[str, Any] = {}
702701
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-
}
702+
replay_properties = _mcp_host_history_fields(
703+
host_payload,
704+
_model_items_for_agui_replay(resolved, llm_str),
705+
)
708706
events.append(
709707
ToolCallResultEvent(
710708
message_id=generate_event_id(),
@@ -1996,10 +1994,11 @@ def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) ->
19961994
"content": llm_result,
19971995
}
19981996
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
1997+
snapshot_message.update(
1998+
_mcp_host_history_fields(
1999+
host_payload,
2000+
_model_items_for_agui_replay(content, llm_result),
2001+
)
20032002
)
20042003
result_by_call_id[call_id] = snapshot_message
20052004
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
@@ -16,15 +16,26 @@
1616
)
1717
from agent_framework._types import ContentType # pyright: ignore[reportPrivateUsage]
1818

19+
from ._state import TOOL_RESULT_DISPLAY_KEY
1920
from ._utils import (
2021
_AGUI_HOST_PAYLOAD_OMITTED_KEY,
2122
_AGUI_MCP_TOOL_RESULT_KEY,
2223
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY,
2324
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
25+
_MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES,
2426
AGUI_TO_FRAMEWORK_ROLE,
2527
FRAMEWORK_TO_AGUI_ROLE,
28+
_bound_host_payload_history,
29+
_extract_mcp_tool_result_host_payload,
30+
_extract_tool_result_marker_values,
31+
_host_payload_history_size,
32+
_mcp_host_history_fields,
2633
_model_content_from_mcp_host_payload,
34+
_model_items_for_agui_replay,
35+
_persistable_host_payload_history,
36+
_project_host_payload_history,
2737
_sanitize_model_replay_item,
38+
_stringify_tool_result,
2839
get_role_value,
2940
normalize_agui_role,
3041
safe_json_parse,
@@ -1039,7 +1050,12 @@ def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, A
10391050
return text, tool_calls
10401051

10411052

1042-
def _split_mixed_message_to_agui(msg: Message, role: str, unresolved_call_ids: set[str]) -> list[dict[str, Any]]:
1053+
def _split_mixed_message_to_agui(
1054+
msg: Message,
1055+
role: str,
1056+
unresolved_call_ids: set[str],
1057+
emitted_results: list[tuple[Content, dict[str, Any]]] | None = None,
1058+
) -> list[dict[str, Any]]:
10431059
"""Convert a Message that carries function_result content into ordered AG-UI messages.
10441060
10451061
A single Agent Framework message can interleave assistant content (text,
@@ -1112,14 +1128,15 @@ def flush_segment() -> None:
11121128
messages.append(assistant_msg)
11131129

11141130
def emit_result(content: Content) -> None:
1115-
messages.append(
1116-
{
1117-
"id": next_id(),
1118-
"role": "tool",
1119-
"content": content.result if content.result is not None else "",
1120-
"toolCallId": content.call_id,
1121-
}
1122-
)
1131+
tool_message: dict[str, Any] = {
1132+
"id": next_id(),
1133+
"role": "tool",
1134+
"content": content.result if content.result is not None else "",
1135+
"toolCallId": content.call_id,
1136+
}
1137+
messages.append(tool_message)
1138+
if emitted_results is not None:
1139+
emitted_results.append((content, tool_message))
11231140
if content.call_id is not None:
11241141
unresolved_call_ids.discard(str(content.call_id))
11251142

@@ -1168,15 +1185,13 @@ def drain_queued() -> None:
11681185
return messages
11691186

11701187

1171-
def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]:
1172-
"""Convert Agent Framework messages to AG-UI format.
1173-
1174-
Args:
1175-
messages: List of Agent Framework Message objects or AG-UI dicts (already converted)
1176-
1177-
Returns:
1178-
List of AG-UI message dictionaries
1179-
"""
1188+
def _convert_agent_framework_messages_to_agui(
1189+
messages: list[Message] | list[dict[str, Any]],
1190+
*,
1191+
emitted_results: list[tuple[Content, dict[str, Any]]] | None = None,
1192+
preserve_host_history_dicts: bool = False,
1193+
) -> list[dict[str, Any]]:
1194+
"""Convert Agent Framework messages to AG-UI format."""
11801195
from ._utils import generate_event_id
11811196

11821197
result: list[dict[str, Any]] = []
@@ -1204,6 +1219,12 @@ def track_emitted(
12041219
if isinstance(msg, dict):
12051220
# Always work on a copy to avoid mutating input
12061221
normalized_msg = msg.copy()
1222+
if not preserve_host_history_dicts and normalized_msg.get(_AGUI_MCP_TOOL_RESULT_KEY) is True:
1223+
normalized_msg = _persistable_host_payload_history([normalized_msg])[0].copy()
1224+
normalized_msg.pop(_AGUI_MCP_TOOL_RESULT_KEY, None)
1225+
normalized_msg.pop(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, None)
1226+
normalized_msg.pop(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, None)
1227+
normalized_msg.pop(_AGUI_HOST_PAYLOAD_OMITTED_KEY, None)
12071228
normalized_msg["role"] = normalize_agui_role(normalized_msg.get("role"))
12081229
# Ensure ID exists
12091230
if "id" not in normalized_msg:
@@ -1235,7 +1256,14 @@ def track_emitted(
12351256
# result is dropped and each result stays after its matching call. Messages
12361257
# with no result use the simple single-message form below.
12371258
if any(content.type == "function_result" for content in msg.contents):
1238-
result.extend(_split_mixed_message_to_agui(msg, role, unresolved_call_ids))
1259+
result.extend(
1260+
_split_mixed_message_to_agui(
1261+
msg,
1262+
role,
1263+
unresolved_call_ids,
1264+
emitted_results,
1265+
)
1266+
)
12391267
continue
12401268

12411269
content_text, tool_calls = _encode_agui_segment(msg.contents)
@@ -1255,6 +1283,83 @@ def track_emitted(
12551283
return result
12561284

12571285

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

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

Lines changed: 5 additions & 12 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] = (
816-
[{"type": "text", "text": result_content}] if model_items is None else model_items
811+
event_replay_properties = _mcp_host_history_fields(
812+
snapshot_result_content,
813+
[{"type": "text", "text": result_content}] if model_items is None else model_items,
817814
)
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-
}
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
@@ -204,6 +204,15 @@ def _host_payload_history_size(message: dict[str, Any]) -> int:
204204
return content_size + sidecar_size
205205

206206

207+
def _mcp_host_history_fields(host_payload: Any, model_items: list[dict[str, Any]]) -> dict[str, Any]:
208+
"""Build the private fields that preserve one MCP Host result for safe replay."""
209+
return {
210+
_AGUI_MCP_TOOL_RESULT_KEY: True,
211+
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: _stringify_tool_result(host_payload),
212+
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: model_items,
213+
}
214+
215+
207216
def _persistable_host_payload_history(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
208217
"""Keep canonical persisted content safe for readers that ignore private replay fields."""
209218
persisted: list[dict[str, Any]] = []

0 commit comments

Comments
 (0)