From f24649dc9e7b7fdf7ab216c887810050bccd39ad Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 8 Sep 2026 11:15:44 +0200 Subject: [PATCH] Python: add MCP Host history conversion for AG-UI Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- python/packages/ag-ui/AGENTS.md | 2 + .../ag-ui/agent_framework_ag_ui/__init__.py | 2 + .../ag-ui/agent_framework_ag_ui/_agent_run.py | 21 +- .../_message_adapters.py | 143 ++++++- .../agent_framework_ag_ui/_run_common.py | 17 +- .../ag-ui/agent_framework_ag_ui/_utils.py | 9 + .../tests/ag_ui/test_message_adapters.py | 382 ++++++++++++++++++ .../core/agent_framework/ag_ui/__init__.py | 2 + .../core/agent_framework/ag_ui/__init__.pyi | 2 + 9 files changed, 538 insertions(+), 42 deletions(-) diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index ac157998dcf..b1ee45f65dd 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -8,6 +8,8 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - **`AgentFrameworkWorkflow`** - Wraps native `Workflow` objects, or accepts `workflow_factory(thread_id)` for thread-scoped workflow instances without subclassing - **`AGUIChatClient`** - Chat client that speaks AG-UI protocol - **`AGUIHttpService`** - HTTP service for AG-UI endpoints +- **`agent_framework_messages_to_agui_host_history()`** - Converts persisted Agent Framework messages to bounded + AG-UI Host history while retaining MCP widget payloads and model replay metadata - **`AGUIEventConverter`** - Converts between Agent Framework and AG-UI events - **`add_agent_framework_fastapi_endpoint()`** - Add AG-UI endpoint to FastAPI app (`SupportsAgentRun` or `Workflow`) - **`InMemoryAGUIThreadSnapshotStore`** - Memory-only latest AG-UI Thread Snapshot store for local development, demos, and tests diff --git a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py index 8df80c25782..622baf89cd6 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py @@ -10,6 +10,7 @@ from ._endpoint import add_agent_framework_fastapi_endpoint from ._event_converters import AGUIEventConverter from ._http_service import AGUIHttpService +from ._message_adapters import agent_framework_messages_to_agui_host_history from ._snapshots import ( DEFAULT_MAX_THREAD_SNAPSHOTS, AGUIThreadID, @@ -39,6 +40,7 @@ "AgentFrameworkWorkflow", "WorkflowFactory", "add_agent_framework_fastapi_endpoint", + "agent_framework_messages_to_agui_host_history", "AGUIChatClient", "AGUIChatOptions", "AGUIEventConverter", diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 7069ae4c415..e7bd836b587 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -103,11 +103,10 @@ from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts from ._utils import ( _AGUI_MCP_TOOL_RESULT_KEY, - _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, - _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, _approval_interrupt_id, _bound_host_payload_history, _function_call_server_label, + _mcp_host_history_fields, _model_items_for_agui_replay, _persistable_host_payload_history, _project_host_payload_history, @@ -700,11 +699,10 @@ def _make_approval_tool_result_events(resolved_approval_results: list[Content]) ui_str = _resolve_ui_payload(llm_str, host_payload if has_host_payload else display_result) replay_properties: dict[str, Any] = {} if has_host_payload: - replay_properties = { - _AGUI_MCP_TOOL_RESULT_KEY: True, - _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: _stringify_tool_result(host_payload), - _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: _model_items_for_agui_replay(resolved, llm_str), - } + replay_properties = _mcp_host_history_fields( + host_payload, + _model_items_for_agui_replay(resolved, llm_str), + ) events.append( ToolCallResultEvent( message_id=generate_event_id(), @@ -1996,10 +1994,11 @@ def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) -> "content": llm_result, } if has_host_payload: - snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True - snapshot_message[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] = _stringify_tool_result(host_payload) - snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = _model_items_for_agui_replay( - content, llm_result + snapshot_message.update( + _mcp_host_history_fields( + host_payload, + _model_items_for_agui_replay(content, llm_result), + ) ) result_by_call_id[call_id] = snapshot_message return result_by_call_id diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index 77751126687..afa1a856616 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -16,15 +16,26 @@ ) from agent_framework._types import ContentType # pyright: ignore[reportPrivateUsage] +from ._state import TOOL_RESULT_DISPLAY_KEY from ._utils import ( _AGUI_HOST_PAYLOAD_OMITTED_KEY, _AGUI_MCP_TOOL_RESULT_KEY, _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + _MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES, AGUI_TO_FRAMEWORK_ROLE, FRAMEWORK_TO_AGUI_ROLE, + _bound_host_payload_history, + _extract_mcp_tool_result_host_payload, + _extract_tool_result_marker_values, + _host_payload_history_size, + _mcp_host_history_fields, _model_content_from_mcp_host_payload, + _model_items_for_agui_replay, + _persistable_host_payload_history, + _project_host_payload_history, _sanitize_model_replay_item, + _stringify_tool_result, get_role_value, normalize_agui_role, safe_json_parse, @@ -1039,7 +1050,12 @@ def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, A return text, tool_calls -def _split_mixed_message_to_agui(msg: Message, role: str, unresolved_call_ids: set[str]) -> list[dict[str, Any]]: +def _split_mixed_message_to_agui( + msg: Message, + role: str, + unresolved_call_ids: set[str], + emitted_results: list[tuple[Content, dict[str, Any]]] | None = None, +) -> list[dict[str, Any]]: """Convert a Message that carries function_result content into ordered AG-UI messages. A single Agent Framework message can interleave assistant content (text, @@ -1112,14 +1128,15 @@ def flush_segment() -> None: messages.append(assistant_msg) def emit_result(content: Content) -> None: - messages.append( - { - "id": next_id(), - "role": "tool", - "content": content.result if content.result is not None else "", - "toolCallId": content.call_id, - } - ) + tool_message: dict[str, Any] = { + "id": next_id(), + "role": "tool", + "content": content.result if content.result is not None else "", + "toolCallId": content.call_id, + } + messages.append(tool_message) + if emitted_results is not None: + emitted_results.append((content, tool_message)) if content.call_id is not None: unresolved_call_ids.discard(str(content.call_id)) @@ -1168,15 +1185,13 @@ def drain_queued() -> None: return messages -def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]: - """Convert Agent Framework messages to AG-UI format. - - Args: - messages: List of Agent Framework Message objects or AG-UI dicts (already converted) - - Returns: - List of AG-UI message dictionaries - """ +def _convert_agent_framework_messages_to_agui( + messages: list[Message] | list[dict[str, Any]], + *, + emitted_results: list[tuple[Content, dict[str, Any]]] | None = None, + preserve_host_history_dicts: bool = False, +) -> list[dict[str, Any]]: + """Convert Agent Framework messages to AG-UI format.""" from ._utils import generate_event_id result: list[dict[str, Any]] = [] @@ -1204,6 +1219,12 @@ def track_emitted( if isinstance(msg, dict): # Always work on a copy to avoid mutating input normalized_msg = msg.copy() + if not preserve_host_history_dicts and normalized_msg.get(_AGUI_MCP_TOOL_RESULT_KEY) is True: + normalized_msg = _persistable_host_payload_history([normalized_msg])[0].copy() + normalized_msg.pop(_AGUI_MCP_TOOL_RESULT_KEY, None) + normalized_msg.pop(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, None) + normalized_msg.pop(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, None) + normalized_msg.pop(_AGUI_HOST_PAYLOAD_OMITTED_KEY, None) normalized_msg["role"] = normalize_agui_role(normalized_msg.get("role")) # Ensure ID exists if "id" not in normalized_msg: @@ -1235,7 +1256,14 @@ def track_emitted( # result is dropped and each result stays after its matching call. Messages # with no result use the simple single-message form below. if any(content.type == "function_result" for content in msg.contents): - result.extend(_split_mixed_message_to_agui(msg, role, unresolved_call_ids)) + result.extend( + _split_mixed_message_to_agui( + msg, + role, + unresolved_call_ids, + emitted_results, + ) + ) continue content_text, tool_calls = _encode_agui_segment(msg.contents) @@ -1255,6 +1283,83 @@ def track_emitted( return result +def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]: + """Convert Agent Framework messages to model-safe AG-UI request format.""" + return _convert_agent_framework_messages_to_agui(messages) + + +def _prepare_host_history_fields( + emitted_results: list[tuple[Content, dict[str, Any]]], + *, + max_size_bytes: int, +) -> tuple[dict[int, dict[str, Any]], set[int]]: + """Materialize only the newest MCP Host projections that fit the aggregate budget.""" + retained_fields: dict[int, dict[str, Any]] = {} + omitted_ids: set[int] = set() + retained_size = 0 + budget_exhausted = False + + for content, _ in reversed(emitted_results): + has_host_payload, host_payload = _extract_mcp_tool_result_host_payload(content) + if not has_host_payload: + continue + content_id = id(content) + if budget_exhausted: + omitted_ids.add(content_id) + continue + + display_values = _extract_tool_result_marker_values(content, TOOL_RESULT_DISPLAY_KEY) + if display_values: + host_payload = display_values[-1] + model_result = _stringify_tool_result(content.result if content.result is not None else "") + fields = _mcp_host_history_fields( + host_payload, + _model_items_for_agui_replay(content, model_result), + ) + message_size = _host_payload_history_size({"content": model_result, **fields}) + if retained_size + message_size > max_size_bytes: + omitted_ids.add(content_id) + budget_exhausted = True + continue + retained_size += message_size + retained_fields[content_id] = fields + + return retained_fields, omitted_ids + + +def agent_framework_messages_to_agui_host_history( + messages: list[Message] | list[dict[str, Any]], + *, + max_host_payload_history_size_bytes: int = _MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES, +) -> list[dict[str, Any]]: + """Convert Agent Framework messages to bounded AG-UI Host history with replay metadata.""" + if messages and isinstance(messages[0], dict): + converted = _persistable_host_payload_history( + _convert_agent_framework_messages_to_agui(messages, preserve_host_history_dicts=True) + ) + else: + message_objects = cast(list[Message], messages) + emitted_results: list[tuple[Content, dict[str, Any]]] = [] + converted = _convert_agent_framework_messages_to_agui( + message_objects, + emitted_results=emitted_results, + ) + host_history_fields, omitted_ids = _prepare_host_history_fields( + emitted_results, + max_size_bytes=max_host_payload_history_size_bytes, + ) + for content, tool_message in emitted_results: + if fields := host_history_fields.get(id(content)): + tool_message.update(fields) + elif id(content) in omitted_ids: + tool_message[_AGUI_HOST_PAYLOAD_OMITTED_KEY] = True + bounded = _bound_host_payload_history( + converted, + max_size_bytes=max_host_payload_history_size_bytes, + ) + return _project_host_payload_history(bounded) + + def extract_text_from_contents(contents: list[Any]) -> str: """Extract text from Agent Framework contents. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 4b8b2e533d2..4f7c33826b1 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -38,12 +38,10 @@ from ._predictive_state import PredictiveStateHandler from ._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY from ._utils import ( - _AGUI_MCP_TOOL_RESULT_KEY, - _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, - _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, _approval_interrupt_id, _extract_mcp_tool_result_host_payload, _extract_tool_result_marker_values, + _mcp_host_history_fields, _model_items_for_agui_replay, _stringify_tool_result, generate_event_id, @@ -810,16 +808,11 @@ def _emit_tool_result_common( } event_replay_properties: dict[str, Any] = {} if snapshot_result is not _UNSET: - snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True - snapshot_message[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] = snapshot_result_content - snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = ( - [{"type": "text", "text": result_content}] if model_items is None else model_items + event_replay_properties = _mcp_host_history_fields( + snapshot_result_content, + [{"type": "text", "text": result_content}] if model_items is None else model_items, ) - event_replay_properties = { - _AGUI_MCP_TOOL_RESULT_KEY: True, - _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: snapshot_result_content, - _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY], - } + snapshot_message.update(event_replay_properties) events[-1] = ToolCallResultEvent( message_id=message_id, tool_call_id=call_id, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py index 770150ccee0..7fe8ae0c832 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py @@ -204,6 +204,15 @@ def _host_payload_history_size(message: dict[str, Any]) -> int: return content_size + sidecar_size +def _mcp_host_history_fields(host_payload: Any, model_items: list[dict[str, Any]]) -> dict[str, Any]: + """Build the private fields that preserve one MCP Host result for safe replay.""" + return { + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: _stringify_tool_result(host_payload), + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: model_items, + } + + def _persistable_host_payload_history(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: """Keep canonical persisted content safe for readers that ignore private replay fields.""" persisted: list[dict[str, Any]] = [] diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index 8b8aa1317c6..e985c083c88 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -6,6 +6,7 @@ import json import logging from itertools import permutations +from types import SimpleNamespace from typing import Any import pytest @@ -13,6 +14,7 @@ from agent_framework_ag_ui._message_adapters import ( agent_framework_messages_to_agui, + agent_framework_messages_to_agui_host_history, agui_messages_to_agent_framework, agui_messages_to_snapshot_format, extract_text_from_contents, @@ -24,6 +26,8 @@ _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, + _host_payload_history_size, + _mcp_host_history_fields, _model_items_for_agui_replay, ) @@ -59,6 +63,384 @@ def test_agent_framework_to_agui_basic(sample_agent_framework_message): assert messages[0]["id"] == "msg-123" +def test_agent_framework_to_agui_preserves_mcp_host_payload_after_reload(): + """Host history uses persisted MCP data without changing generic model output.""" + host_payload = { + "content": [{"type": "text", "text": "Summary"}], + "structuredContent": {"image_url": "https://example.test/widget.png"}, + "isError": False, + } + tool_return = Content.from_text( + "Summary", + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + message = Message( + role="tool", + contents=[Content.from_function_result(call_id="mcp-1", result=[tool_return])], + message_id="message-1", + ) + restored_message = Message.from_dict(message.to_dict()) + + outbound = agent_framework_messages_to_agui([restored_message]) + converted = agent_framework_messages_to_agui_host_history([restored_message]) + + assert restored_message.contents[0].result == "Summary" + assert outbound[0]["content"] == "Summary" + assert _AGUI_MCP_TOOL_RESULT_KEY not in outbound[0] + assert _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY not in outbound[0] + assert _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY not in outbound[0] + assert json.loads(converted[0]["content"]) == host_payload + assert converted[0]["toolCallId"] == "mcp-1" + assert _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY not in converted[0] + + provider_messages = agui_messages_to_agent_framework(converted) + assert provider_messages[0].contents[0].result == "Summary" + + +def test_generic_conversion_strips_host_history_from_dict_input(): + """Generic outbound conversion reduces Host-history dictionaries to model-safe content.""" + host_history = { + "id": "mcp-result", + "role": "tool", + "toolCallId": "mcp-call", + "content": json.dumps({"content": [{"type": "text", "text": "Host only"}]}), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": "Model summary"}], + } + + converted = agent_framework_messages_to_agui([host_history]) + + assert converted[0]["content"] == "Model summary" + assert _AGUI_MCP_TOOL_RESULT_KEY not in converted[0] + assert _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY not in converted[0] + assert _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY not in converted[0] + + +def test_host_history_converter_reprojects_projected_and_persisted_dicts(): + """Both public dictionary forms reproject Host data and retain model-safe replay.""" + projected_host = { + "accepted": True, + "content": [{"type": "text", "text": "Host approval lookalike"}], + "structuredContent": {"source": "projected"}, + } + persisted_host = { + "content": [{"type": "text", "text": "Server-only text"}], + "structuredContent": {"source": "persisted"}, + "isError": False, + } + projected = { + "id": "projected", + "role": "tool", + "toolCallId": "projected-call", + "content": json.dumps(projected_host), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": "Model-safe result"}], + } + persisted = { + "id": "persisted", + "role": "tool", + "toolCallId": "persisted-call", + "content": "", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: persisted_host, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [], + } + + converted = agent_framework_messages_to_agui_host_history([projected, persisted]) + + assert json.loads(converted[0]["content"]) == projected_host + assert converted[0][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [{"type": "text", "text": "Model-safe result"}] + assert json.loads(converted[1]["content"]) == persisted_host + assert converted[1][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [] + assert all(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY not in message for message in converted) + + replayed = agui_messages_to_agent_framework(converted) + assert replayed[0].contents[0].type == "function_result" + assert replayed[0].contents[0].result == "Model-safe result" + assert replayed[1].contents[0].result == "" + assert replayed[1].contents[0].items == [] + + +def test_host_history_converter_preserves_malformed_sidecar_for_safe_replay(): + """Malformed persisted sidecars remain marked so inbound replay uses the safe Host fallback.""" + host_payload = { + "content": [{"type": "text", "text": "Safe Host summary"}], + "structuredContent": {"private": True}, + "isError": False, + } + persisted = { + "role": "tool", + "toolCallId": "malformed-call", + "content": "Tool result unavailable.", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": {"invalid": True}}], + } + + converted = agent_framework_messages_to_agui_host_history([persisted]) + + assert json.loads(converted[0]["content"]) == host_payload + assert converted[0][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == persisted[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] + replayed = agui_messages_to_agent_framework(converted) + assert replayed[0].contents[0].result == "Safe Host summary" + + +def test_host_history_converter_bounds_non_string_persisted_host_values(): + """Non-string Host values count toward the same aggregate Host-plus-sidecar budget.""" + persisted_messages = [ + { + "id": f"result-{index}", + "role": "tool", + "toolCallId": f"call-{index}", + "content": f"Model {index}", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": f"Model {index}"}], + } + for index, host_payload in enumerate( + [ + {"structuredContent": {"index": 0, "data": "x" * 80}}, + [{"type": "resource", "resource": {"text": "y" * 80}}], + ] + ) + ] + unbounded = agent_framework_messages_to_agui_host_history( + persisted_messages, + max_host_payload_history_size_bytes=10_000, + ) + newest_size = _host_payload_history_size(unbounded[1]) + + converted = agent_framework_messages_to_agui_host_history( + persisted_messages, + max_host_payload_history_size_bytes=newest_size, + ) + + assert converted[0]["content"] == "Model 0" + assert converted[0][_AGUI_HOST_PAYLOAD_OMITTED_KEY] is True + assert _AGUI_MCP_TOOL_RESULT_KEY not in converted[0] + assert json.loads(converted[1]["content"]) == persisted_messages[1][_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] + assert converted[1][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [{"type": "text", "text": "Model 1"}] + + +def test_host_history_converter_makes_model_replay_metadata_json_safe(): + """Provider-visible model metadata is JSON-safe in public Host history.""" + host_payload = {"content": [{"type": "text", "text": "Host result"}], "isError": False} + model_item = Content.from_text( + "Model result", + additional_properties={ + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload, + "provider_visible": SimpleNamespace(value="kept"), + }, + ) + message = Message( + role="tool", + contents=[Content.from_function_result(call_id="json-safe", result=[model_item])], + ) + + converted = agent_framework_messages_to_agui_host_history([message]) + + assert converted[0][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY][0]["additional_properties"]["provider_visible"] == { + "value": "kept" + } + json.dumps(converted) + + +def test_host_history_conversion_preserves_parallel_results_and_mixed_content(): + """Host conversion keeps upstream parallel-result splitting and mixed content.""" + host_payload = { + "content": [{"type": "text", "text": "Host summary"}], + "structuredContent": {"widget": "parallel"}, + "isError": False, + } + mcp_result = Content.from_text( + "Model summary", + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + message = Message( + role="assistant", + contents=[ + Content.from_function_result(call_id="mcp-call", result=[mcp_result]), + Content.from_function_result(call_id="plain-call", result="Plain result"), + Content.from_text("Both tools completed."), + ], + message_id="parallel-result", + ) + + generic = agent_framework_messages_to_agui([message]) + host_history = agent_framework_messages_to_agui_host_history([message]) + + assert [item["role"] for item in generic] == ["tool", "tool", "assistant"] + assert [item["toolCallId"] for item in generic[:2]] == ["mcp-call", "plain-call"] + assert [item["content"] for item in generic] == ["Model summary", "Plain result", "Both tools completed."] + assert all(_AGUI_MCP_TOOL_RESULT_KEY not in item for item in generic) + + assert [item["role"] for item in host_history] == ["tool", "tool", "assistant"] + assert [item["toolCallId"] for item in host_history[:2]] == ["mcp-call", "plain-call"] + assert json.loads(host_history[0]["content"]) == host_payload + assert host_history[0][_AGUI_MCP_TOOL_RESULT_KEY] is True + assert host_history[1]["content"] == "Plain result" + assert host_history[2]["content"] == "Both tools completed." + assert len({item["id"] for item in host_history}) == 3 + + +def test_host_history_conversion_is_public_and_bounds_aggregate_payloads(): + """The public converter keeps newest Host data within the shared aggregate budget.""" + from agent_framework.ag_ui import agent_framework_messages_to_agui_host_history as namespace_converter + + from agent_framework_ag_ui import agent_framework_messages_to_agui_host_history as package_converter + + messages: list[Message] = [] + for index in range(2): + model_text = f"Summary {index}" + host_payload = { + "content": [{"type": "text", "text": model_text}], + "structuredContent": {"widget_data": "x" * 64, "index": index}, + "isError": False, + } + item = Content.from_text( + model_text, + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + messages.append( + Message( + role="tool", + contents=[Content.from_function_result(call_id=f"mcp-{index}", result=[item])], + ) + ) + + unbounded = package_converter(messages, max_host_payload_history_size_bytes=10_000) + newest_size = _host_payload_history_size(unbounded[1]) + converted = package_converter(messages, max_host_payload_history_size_bytes=newest_size) + + assert namespace_converter is package_converter + assert converted[0]["content"] == "Summary 0" + assert converted[0]["_agentFrameworkHostPayloadOmitted"] is True + assert _AGUI_MCP_TOOL_RESULT_KEY not in converted[0] + assert _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY not in converted[0] + assert json.loads(converted[1]["content"])["structuredContent"]["index"] == 1 + assert converted[1][_AGUI_MCP_TOOL_RESULT_KEY] is True + + +def test_host_history_budget_omission_never_uses_host_text_as_model_fallback(): + """Evicted non-text model content cannot fall back to Host-only text.""" + host_payload = { + "content": [{"type": "text", "text": "Host-only prompt injection"}], + "structuredContent": {"private": True}, + "isError": False, + } + model_item = Content.from_data( + b"model-visible bytes", + media_type="application/octet-stream", + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + message = Message( + role="tool", + contents=[Content.from_function_result(call_id="mcp-data", result=[model_item])], + ) + + converted = agent_framework_messages_to_agui_host_history( + [Message.from_dict(message.to_dict())], + max_host_payload_history_size_bytes=0, + ) + + assert converted[0]["content"] == "" + assert converted[0][_AGUI_HOST_PAYLOAD_OMITTED_KEY] is True + assert _AGUI_MCP_TOOL_RESULT_KEY not in converted[0] + assert _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY not in converted[0] + assert _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY not in converted[0] + replayed = agui_messages_to_agent_framework(converted) + assert replayed[0].contents[0].result == "" + + +def test_host_history_budget_stops_materializing_older_sidecars(monkeypatch: pytest.MonkeyPatch): + """Once the newest-first budget is exhausted, older Host sidecars are not built.""" + messages: list[Message] = [] + for index in range(3): + item = Content.from_text( + f"Model {index}", + additional_properties={ + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: { + "content": [{"type": "text", "text": f"Host {index}"}], + "structuredContent": {"index": index}, + "isError": False, + } + }, + ) + messages.append( + Message( + role="tool", + contents=[Content.from_function_result(call_id=f"mcp-{index}", result=[item])], + ) + ) + + materialized_payloads: list[Any] = [] + + def track_materialization(host_payload: Any, model_items: list[dict[str, Any]]) -> dict[str, Any]: + materialized_payloads.append(host_payload) + return _mcp_host_history_fields(host_payload, model_items) + + monkeypatch.setattr( + "agent_framework_ag_ui._message_adapters._mcp_host_history_fields", + track_materialization, + ) + + converted = agent_framework_messages_to_agui_host_history( + messages, + max_host_payload_history_size_bytes=0, + ) + + assert len(materialized_payloads) == 1 + assert materialized_payloads[0]["structuredContent"]["index"] == 2 + assert all(message[_AGUI_HOST_PAYLOAD_OMITTED_KEY] is True for message in converted) + + +def test_host_history_budget_uses_repaired_parallel_result_order(): + """Newest-first retention follows final AG-UI order after buffered results are repaired.""" + + def host_result(call_id: str) -> Content: + item = Content.from_text( + f"Model {call_id}", + additional_properties={ + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: { + "content": [{"type": "text", "text": f"Host {call_id}"}], + "structuredContent": {"call_id": call_id}, + "isError": False, + } + }, + ) + return Content.from_function_result(call_id=call_id, result=[item]) + + message = Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="a", name="tool_a", arguments={}), + Content.from_function_call(call_id="b", name="tool_b", arguments={}), + host_result("a"), + Content.from_function_call(call_id="c", name="tool_c", arguments={}), + host_result("c"), + host_result("b"), + ], + ) + unbounded = agent_framework_messages_to_agui_host_history( + [message], + max_host_payload_history_size_bytes=10_000, + ) + newest = next(item for item in unbounded if item.get("toolCallId") == "c") + newest_size = _host_payload_history_size(newest) + + converted = agent_framework_messages_to_agui_host_history( + [message], + max_host_payload_history_size_bytes=newest_size, + ) + tool_messages = [item for item in converted if item["role"] == "tool"] + + assert [item["toolCallId"] for item in tool_messages] == ["a", "b", "c"] + assert tool_messages[0][_AGUI_HOST_PAYLOAD_OMITTED_KEY] is True + assert tool_messages[1][_AGUI_HOST_PAYLOAD_OMITTED_KEY] is True + assert json.loads(tool_messages[2]["content"])["structuredContent"]["call_id"] == "c" + assert tool_messages[2][_AGUI_MCP_TOOL_RESULT_KEY] is True + + def test_marked_mcp_snapshot_restores_lossless_model_items(): """Inbound replay restores media and provider-visible data but excludes Host-only metadata.""" host_payload = { diff --git a/python/packages/core/agent_framework/ag_ui/__init__.py b/python/packages/core/agent_framework/ag_ui/__init__.py index 580ae153a9a..f5a3e1986f1 100644 --- a/python/packages/core/agent_framework/ag_ui/__init__.py +++ b/python/packages/core/agent_framework/ag_ui/__init__.py @@ -16,6 +16,7 @@ - InMemoryAGUIThreadSnapshotStore - SnapshotScopeResolver - add_agent_framework_fastapi_endpoint +- agent_framework_messages_to_agui_host_history - state_update - __version__ """ @@ -29,6 +30,7 @@ "AgentFrameworkAgent", "AgentFrameworkWorkflow", "add_agent_framework_fastapi_endpoint", + "agent_framework_messages_to_agui_host_history", "AGUIChatClient", "AGUIEventConverter", "AGUIHttpService", diff --git a/python/packages/core/agent_framework/ag_ui/__init__.pyi b/python/packages/core/agent_framework/ag_ui/__init__.pyi index e57ba45ac62..ebf00eb0663 100644 --- a/python/packages/core/agent_framework/ag_ui/__init__.pyi +++ b/python/packages/core/agent_framework/ag_ui/__init__.pyi @@ -12,6 +12,7 @@ from agent_framework_ag_ui import ( SnapshotScopeResolver, __version__, add_agent_framework_fastapi_endpoint, + agent_framework_messages_to_agui_host_history, state_update, ) @@ -27,5 +28,6 @@ __all__ = [ "SnapshotScopeResolver", "__version__", "add_agent_framework_fastapi_endpoint", + "agent_framework_messages_to_agui_host_history", "state_update", ]