Skip to content

Commit c5d5f82

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

9 files changed

Lines changed: 539 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 & 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: 125 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -8,24 +8,34 @@
88
import binascii
99
import json
1010
import logging
11-
from typing import Any, cast
11+
from typing import Any, cast, get_args
1212

1313
from agent_framework import (
1414
Content,
1515
Message,
1616
)
1717
from agent_framework._types import ContentType # pyright: ignore[reportPrivateUsage]
18-
from typing_extensions import get_args
1918

19+
from ._state import TOOL_RESULT_DISPLAY_KEY
2020
from ._utils import (
2121
_AGUI_HOST_PAYLOAD_OMITTED_KEY,
2222
_AGUI_MCP_TOOL_RESULT_KEY,
2323
_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY,
2424
_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY,
25+
_MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES,
2526
AGUI_TO_FRAMEWORK_ROLE,
2627
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,
2733
_model_content_from_mcp_host_payload,
34+
_model_items_for_agui_replay,
35+
_persistable_host_payload_history,
36+
_project_host_payload_history,
2837
_sanitize_model_replay_item,
38+
_stringify_tool_result,
2939
get_role_value,
3040
normalize_agui_role,
3141
safe_json_parse,
@@ -1040,7 +1050,12 @@ def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, A
10401050
return text, tool_calls
10411051

10421052

1043-
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]]:
10441059
"""Convert a Message that carries function_result content into ordered AG-UI messages.
10451060
10461061
A single Agent Framework message can interleave assistant content (text,
@@ -1113,14 +1128,15 @@ def flush_segment() -> None:
11131128
messages.append(assistant_msg)
11141129

11151130
def emit_result(content: Content) -> None:
1116-
messages.append(
1117-
{
1118-
"id": next_id(),
1119-
"role": "tool",
1120-
"content": content.result if content.result is not None else "",
1121-
"toolCallId": content.call_id,
1122-
}
1123-
)
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))
11241140
if content.call_id is not None:
11251141
unresolved_call_ids.discard(str(content.call_id))
11261142

@@ -1169,15 +1185,13 @@ def drain_queued() -> None:
11691185
return messages
11701186

11711187

1172-
def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]:
1173-
"""Convert Agent Framework messages to AG-UI format.
1174-
1175-
Args:
1176-
messages: List of Agent Framework Message objects or AG-UI dicts (already converted)
1177-
1178-
Returns:
1179-
List of AG-UI message dictionaries
1180-
"""
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."""
11811195
from ._utils import generate_event_id
11821196

11831197
result: list[dict[str, Any]] = []
@@ -1205,6 +1219,12 @@ def track_emitted(
12051219
if isinstance(msg, dict):
12061220
# Always work on a copy to avoid mutating input
12071221
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)
12081228
normalized_msg["role"] = normalize_agui_role(normalized_msg.get("role"))
12091229
# Ensure ID exists
12101230
if "id" not in normalized_msg:
@@ -1236,7 +1256,14 @@ def track_emitted(
12361256
# result is dropped and each result stays after its matching call. Messages
12371257
# with no result use the simple single-message form below.
12381258
if any(content.type == "function_result" for content in msg.contents):
1239-
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+
)
12401267
continue
12411268

12421269
content_text, tool_calls = _encode_agui_segment(msg.contents)
@@ -1256,6 +1283,83 @@ def track_emitted(
12561283
return result
12571284

12581285

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+
12591363
def extract_text_from_contents(contents: list[Any]) -> str:
12601364
"""Extract text from Agent Framework contents.
12611365

‎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)