From 58f55702b4f51f0613d2fd335f1d217bffecc76d Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 7 Sep 2026 20:55:17 +0200 Subject: [PATCH] Python: preserve MCP Host payloads in AG-UI snapshots Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 56 ++- .../_message_adapters.py | 56 ++- .../agent_framework_ag_ui/_run_common.py | 91 +++-- .../_snapshot_session.py | 4 +- .../ag-ui/agent_framework_ag_ui/_utils.py | 225 ++++++++++- .../ag-ui/agent_framework_ag_ui/_workflow.py | 42 +- .../tests/ag_ui/test_message_adapters.py | 282 +++++++++++++ .../ag-ui/tests/ag_ui/test_run_common.py | 377 +++++++++++++++++- .../tests/ag_ui/test_snapshot_session.py | 32 ++ .../ag-ui/tests/ag_ui/test_snapshots.py | 33 ++ python/packages/core/agent_framework/_mcp.py | 54 ++- .../packages/core/agent_framework/_tools.py | 32 +- .../core/test_function_invocation_logic.py | 49 ++- python/packages/core/tests/core/test_mcp.py | 189 ++++++++- 14 files changed, 1427 insertions(+), 95 deletions(-) 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 4fdaef188d5..7069ae4c415 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 @@ -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 ( @@ -102,8 +102,16 @@ ) 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, + _model_items_for_agui_replay, + _persistable_host_payload_history, + _project_host_payload_history, + _stringify_tool_result, canonical_function_arguments, convert_agui_tools_to_agent_framework, generate_event_id, @@ -687,13 +695,23 @@ 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) + 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), + } events.append( ToolCallResultEvent( message_id=generate_event_id(), tool_call_id=resolved.call_id, content=ui_str, role="tool", + **replay_properties, ) ) return events @@ -1968,12 +1986,22 @@ 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": 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 + ) + result_by_call_id[call_id] = snapshot_message return result_by_call_id @@ -1987,6 +2015,16 @@ def _merge_resolved_approval_results_into_snapshot( snapshot_messages[:] = [message for message in snapshot_messages if not message.get("function_approvals")] return + for message in snapshot_messages: + if normalize_agui_role(message.get("role", "")) != "tool": + continue + tool_call_id = message.get("toolCallId") or message.get("tool_call_id") + if not tool_call_id or message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: + continue + replacement = result_by_call_id.get(str(tool_call_id)) + if replacement is not None and replacement.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: + result_by_call_id.pop(str(tool_call_id)) + merged_messages: list[dict[str, Any]] = [] for message in snapshot_messages: if message.get("function_approvals"): @@ -2067,7 +2105,8 @@ def _build_messages_snapshot( if flow.snapshot_segments: _append_segmented_snapshot_messages(flow, all_messages) - return MessagesSnapshotEvent(messages=all_messages) # type: ignore[arg-type] + bounded_messages = _bound_host_payload_history(_persistable_host_payload_history(all_messages)) + return MessagesSnapshotEvent(messages=_project_host_payload_history(bounded_messages)) # type: ignore[arg-type] # Add assistant message with tool calls only (no content) if flow.pending_tool_calls: @@ -2101,7 +2140,8 @@ 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] + bounded_messages = _bound_host_payload_history(_persistable_host_payload_history(all_messages)) + return MessagesSnapshotEvent(messages=_project_host_payload_history(bounded_messages)) # type: ignore[arg-type] def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str, Any]]: @@ -2876,7 +2916,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(_persistable_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( @@ -3241,7 +3281,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(_persistable_host_payload_history(persisted_messages)), state=latest_state_snapshot, interrupt=flow.interrupts or None, session_state=_safe_serialize_session_continuation_state( 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 24624517ebc..77751126687 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 @@ -8,22 +8,30 @@ import binascii import json import logging -from typing import Any, cast +from typing import Any, cast, get_args from agent_framework import ( Content, Message, ) +from agent_framework._types import ContentType # pyright: ignore[reportPrivateUsage] 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, AGUI_TO_FRAMEWORK_ROLE, FRAMEWORK_TO_AGUI_ROLE, + _model_content_from_mcp_host_payload, + _sanitize_model_replay_item, get_role_value, normalize_agui_role, safe_json_parse, ) logger = logging.getLogger(__name__) +_VALID_CONTENT_TYPES = frozenset(get_args(ContentType)) def _append_synthetic_tool_results( @@ -719,6 +727,52 @@ 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: + host_payload = msg.get(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, result_content) + parsed_host_payload = safe_json_parse(host_payload) + serialized_items = msg.get(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) + function_result: Content | None = None + if isinstance(serialized_items, list) and all( + isinstance(item, dict) and item.get("type") in _VALID_CONTENT_TYPES for item in serialized_items + ): + try: + model_items = [ + Content.from_dict(_sanitize_model_replay_item(item)) for item in serialized_items + ] + if any( + item.get("type") == "text" and "text" in item and not isinstance(item.get("text"), str) + for item in serialized_items + ): + raise TypeError("Serialized text replay content must contain a string") + function_result = Content.from_function_result( + call_id=str(tool_call_id), + result=model_items, + ) + except (RecursionError, TypeError, ValueError): + function_result = None + if function_result is None: + model_items = [Content.from_text(_model_content_from_mcp_host_payload(parsed_host_payload))] + function_result = Content.from_function_result(call_id=str(tool_call_id), result=model_items) + chat_msg = Message( + role="tool", + contents=[function_result], + ) + if "id" in msg: + chat_msg.message_id = msg["id"] + result.append(chat_msg) + continue + + if msg.get(_AGUI_HOST_PAYLOAD_OMITTED_KEY) is True: + safe_result = result_content if isinstance(result_content, (str, dict, list)) else str(result_content) + chat_msg = Message( + role="tool", + contents=[Content.from_function_result(call_id=str(tool_call_id), result=safe_result)], + ) + 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: 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 5f73c58d53b..4b8b2e533d2 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 @@ -37,7 +37,19 @@ 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_HOST_PAYLOAD_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__) @@ -696,22 +708,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``. @@ -745,15 +741,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, @@ -762,6 +762,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. @@ -789,6 +791,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( @@ -799,14 +802,32 @@ 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": result_content, + } + 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 = { + _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], } - ) + events[-1] = ToolCallResultEvent( + message_id=message_id, + tool_call_id=call_id, + content=ui_result_content, + role="tool", + **event_replay_properties, + ) + 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). @@ -851,6 +872,9 @@ 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, @@ -858,6 +882,10 @@ def _emit_tool_result( 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 + ), ) @@ -1022,6 +1050,9 @@ 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, @@ -1029,6 +1060,10 @@ def _emit_mcp_tool_result( 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 + ), ) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py index b0694596e74..cb581d5d4ac 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py @@ -28,7 +28,7 @@ AGUIThreadSnapshotStore, _clear_thread_snapshot_interrupt, ) -from ._utils import make_json_safe +from ._utils import _project_host_payload_history, make_json_safe logger = logging.getLogger(__name__) @@ -105,7 +105,7 @@ async def hydrate_events(self, *, run_id: str) -> AsyncGenerator[BaseEvent]: if snapshot.state is not None: yield StateSnapshotEvent(snapshot=snapshot.state) if snapshot.messages: - yield MessagesSnapshotEvent(messages=snapshot.messages) # type: ignore[arg-type] + yield MessagesSnapshotEvent(messages=_project_host_payload_history(snapshot.messages)) # type: ignore[arg-type] yield _build_run_finished_event(run_id=run_id, thread_id=self._thread_id, interrupts=snapshot.interrupt) def effective_state( 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 62d4f39c26f..770150ccee0 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py @@ -10,9 +10,33 @@ from collections.abc import Callable, MutableMapping, Sequence from typing import Any -from agent_framework import AgentResponseUpdate, ChatResponseUpdate, FunctionTool +from agent_framework import AgentResponseUpdate, ChatResponseUpdate, Content, FunctionTool +from agent_framework import _mcp as _core_mcp # pyright: ignore[reportPrivateUsage] from agent_framework._serialization import make_json_safe # pyright: ignore[reportPrivateUsage] + +def _mcp_tool_result_host_payload_key(core_mcp: Any) -> str: + """Resolve the private core marker while supporting older core packages.""" + return getattr(core_mcp, "_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY", "_mcp_tool_result_host_payload") + + +_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY = _mcp_tool_result_host_payload_key(_core_mcp) +_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY = "_agentFrameworkModelContent" +_AGUI_MCP_TOOL_RESULT_KEY = "_agentFrameworkMcpResult" +_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY = "_agentFrameworkHostPayload" +_AGUI_HOST_PAYLOAD_OMITTED_KEY = "_agentFrameworkHostPayloadOmitted" +_MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES = 8 * 1024 * 1024 +_HOST_ONLY_REPLAY_ITEM_KEYS = frozenset( + { + "_meta", + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + _AGUI_HOST_PAYLOAD_OMITTED_KEY, + } +) + # Role mapping constants AGUI_TO_FRAMEWORK_ROLE: dict[str, str] = { "user": "user", @@ -55,6 +79,205 @@ def safe_json_parse(value: Any) -> dict[str, Any] | None: return None +def _extract_tool_result_marker_values(content: Any, key: str) -> list[Any]: + """Extract marker values from outer and inner tool-result content.""" + values: list[Any] = [] + + outer_properties = getattr(content, "additional_properties", None) or {} + if key in outer_properties: + values.append(outer_properties[key]) + + for item in getattr(content, "items", None) or (): + item_properties = getattr(item, "additional_properties", None) or {} + if key in item_properties: + values.append(item_properties[key]) + + return values + + +def _extract_mcp_tool_result_host_payload(content: Any) -> tuple[bool, Any]: + """Return whether a core-preserved MCP Host payload exists and its value.""" + outer_properties = getattr(content, "additional_properties", None) or {} + if _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY in outer_properties: + return True, outer_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] + + values: list[Any] = [] + for item in getattr(content, "items", None) or (): + item_properties = getattr(item, "additional_properties", None) or {} + if _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY in item_properties: + values.append(item_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]) + return (True, values[-1]) if values else (False, None) + + +def _model_content_from_mcp_host_payload(payload: Any) -> str: + """Recover safe content-only text when marked history lacks a valid sidecar.""" + if not isinstance(payload, dict): + return "Tool result unavailable." + if payload.get("isError") is True: + return "Error: Function failed." + content = payload.get("content") + if not isinstance(content, list): + return "Tool result unavailable." + + text_parts: list[str] = [] + for item in content: + if not isinstance(item, dict): + continue + if item.get("type") == "text" and isinstance(item.get("text"), str): + text_parts.append(item["text"]) + continue + resource = item.get("resource") + if item.get("type") == "resource" and isinstance(resource, dict) and isinstance(resource.get("text"), str): + text_parts.append(resource["text"]) + return "\n".join(text_parts) if text_parts else "null" + + +def _model_items_for_agui_replay(content: Any, model_result: str) -> list[dict[str, Any]]: + """Serialize model-facing items without Host-only MCP metadata.""" + items = getattr(content, "items", None) + if items is None: + output = getattr(content, "output", None) + if isinstance(output, list) and all(isinstance(item, Content) for item in output): + items = output + if items is None: + return [{"type": "text", "text": model_result}] + + try: + serialized_items: list[dict[str, Any]] = [] + for item in items: + serialized_item = make_json_safe(_sanitize_model_replay_item(item.to_dict())) + if not isinstance(serialized_item, dict): + raise TypeError("Serialized model replay item must be a dictionary") + serialized_items.append(serialized_item) + return serialized_items + except (RecursionError, TypeError, ValueError): + return [{"type": "text", "text": model_result}] + + +def _sanitize_model_replay_item(item: dict[str, Any]) -> dict[str, Any]: + """Remove Host-only item metadata without mutating caller-owned replay data.""" + sanitized_item = item.copy() + additional_properties = item.get("additional_properties") + if not isinstance(additional_properties, dict): + return sanitized_item + sanitized_properties = { + key: value for key, value in additional_properties.items() if key not in _HOST_ONLY_REPLAY_ITEM_KEYS + } + if sanitized_properties: + sanitized_item["additional_properties"] = sanitized_properties + else: + sanitized_item.pop("additional_properties", None) + return sanitized_item + + +def _model_text_from_replay_items(message: dict[str, Any]) -> str: + """Return the text projection used when a Host payload is omitted.""" + serialized_items = message.get(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) + if not isinstance(serialized_items, list): + return "Tool result unavailable." + if not serialized_items: + return "" + model_text = "\n".join( + item["text"] + for item in serialized_items + if isinstance(item, dict) and item.get("type") == "text" and isinstance(item.get("text"), str) + ) + return model_text or "Tool result unavailable." + + +def _host_payload_history_size(message: dict[str, Any]) -> int: + """Return aggregate bytes retained for one Host projection and replay sidecar.""" + content = message.get(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, message.get("content")) + serialized_items = message.get(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) + try: + content_size = len(json.dumps(make_json_safe(content), separators=(",", ":")).encode("utf-8")) + except (RecursionError, TypeError, ValueError): + content_size = _MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES + 1 + try: + sidecar_size = ( + len(json.dumps(make_json_safe(serialized_items), separators=(",", ":")).encode("utf-8")) + if isinstance(serialized_items, list) + else 0 + ) + except (RecursionError, TypeError, ValueError): + sidecar_size = _MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES + 1 + return content_size + sidecar_size + + +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]] = [] + for message in messages: + if message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: + persisted.append(message) + continue + persisted_message = message.copy() + host_payload = message.get(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, message.get("content")) + persisted_message["content"] = _model_text_from_replay_items(message) + persisted_message[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] = _stringify_tool_result(host_payload) + persisted.append(persisted_message) + return persisted + + +def _project_host_payload_history(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Project private complete MCP results into AG-UI Host-visible tool content.""" + projected: list[dict[str, Any]] = [] + for message in messages: + host_payload = message.get(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY) + if message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True or host_payload is None: + projected.append(message) + continue + projected_message = message.copy() + projected_message["content"] = _stringify_tool_result(host_payload) + projected_message.pop(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, None) + projected.append(projected_message) + return projected + + +def _bound_host_payload_history( + messages: list[dict[str, Any]], + *, + max_size_bytes: int | None = None, +) -> list[dict[str, Any]]: + """Retain the newest Host projections and sidecars within one fixed budget.""" + if max_size_bytes is None: + max_size_bytes = _MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES + + retained_size = 0 + omit_indices: set[int] = set() + for index in range(len(messages) - 1, -1, -1): + message = messages[index] + if message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: + continue + message_size = _host_payload_history_size(message) + if retained_size + message_size > max_size_bytes: + omit_indices.add(index) + else: + retained_size += message_size + + if not omit_indices: + return messages + + bounded_messages: list[dict[str, Any]] = [] + for index, message in enumerate(messages): + if index not in omit_indices: + bounded_messages.append(message) + continue + bounded_message = message.copy() + bounded_message["content"] = _model_text_from_replay_items(message) + bounded_message.pop(_AGUI_MCP_TOOL_RESULT_KEY, None) + bounded_message.pop(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, None) + bounded_message.pop(_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, None) + bounded_message[_AGUI_HOST_PAYLOAD_OMITTED_KEY] = True + bounded_messages.append(bounded_message) + return bounded_messages + + +def _stringify_tool_result(raw_result: Any) -> str: + """Serialize a tool result for an AG-UI tool message.""" + return raw_result if isinstance(raw_result, str) else json.dumps(make_json_safe(raw_result)) + + def canonical_function_arguments(function_call: Any) -> str | None: """Return a stable representation of function-call arguments.""" if function_call is None: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index 73b2a1d7ec0..4a20e171225 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -41,7 +41,15 @@ AGUIThreadSnapshot, AGUIThreadSnapshotStore, ) -from ._utils import generate_event_id, make_json_safe +from ._utils import ( + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + _bound_host_payload_history, + _persistable_host_payload_history, + generate_event_id, + make_json_safe, +) from ._workflow_run import _pending_request_events, run_workflow_stream # pyright: ignore[reportPrivateUsage] logger = logging.getLogger(__name__) @@ -170,7 +178,12 @@ def build(self) -> AGUIThreadSnapshot: """Return the replayable thread snapshot.""" self._flush_open_text_message() messages = self._emitted_messages if self._emitted_messages is not None else self._synthesized_messages - return AGUIThreadSnapshot(messages=messages, state=self.state, interrupt=self.interrupt) + persisted = _persistable_host_payload_history(messages) + return AGUIThreadSnapshot( + messages=_bound_host_payload_history(persisted), + state=self.state, + interrupt=self.interrupt, + ) def _observe_text_start(self, event: TextMessageStartEvent) -> None: if self._open_text_message is not None and self._open_text_message.get("id") != event.message_id: @@ -223,14 +236,23 @@ def _observe_tool_call_args(self, event: ToolCallArgsEvent) -> None: function_payload["arguments"] = f"{function_payload.get('arguments', '')}{event.delta}" def _observe_tool_call_result(self, event: ToolCallResultEvent) -> None: - self._synthesized_messages.append( - { - "id": event.message_id, - "role": "tool", - "toolCallId": event.tool_call_id, - "content": event.content, - } - ) + message: dict[str, Any] = { + "id": event.message_id, + "role": "tool", + "toolCallId": event.tool_call_id, + "content": event.content, + } + if getattr(event, _AGUI_MCP_TOOL_RESULT_KEY, False) is True: + message[_AGUI_MCP_TOOL_RESULT_KEY] = True + message[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] = getattr( + event, _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, event.content + ) + message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = getattr( + event, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + [{"type": "text", "text": "Tool result unavailable."}], + ) + self._synthesized_messages.append(message) # A result closes the current tool-call group; later tool calls start a new # assistant message so replayed transcripts keep results adjacent to their # tool_calls message, which provider APIs require. 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 a756009df6a..8b8aa1317c6 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 @@ -18,6 +18,14 @@ extract_text_from_contents, normalize_agui_input_messages, ) +from agent_framework_ag_ui._utils import ( + _AGUI_HOST_PAYLOAD_OMITTED_KEY, + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, + _model_items_for_agui_replay, +) @pytest.fixture @@ -51,6 +59,280 @@ def test_agent_framework_to_agui_basic(sample_agent_framework_message): assert messages[0]["id"] == "msg-123" +def test_marked_mcp_snapshot_restores_lossless_model_items(): + """Inbound replay restores media and provider-visible data but excludes Host-only metadata.""" + host_payload = { + "content": [{"type": "image", "data": "aW1hZ2U=", "mimeType": "image/png"}], + "structuredContent": {"widget": "image"}, + "isError": False, + } + model_items = [ + Content.from_text( + "Image ready", + additional_properties={ + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload, + "_meta": {"server_only": True}, + "provider_visible": "kept", + }, + ), + Content.from_data(b"image", media_type="image/png"), + Content.from_uri("https://example.test/resource.txt", media_type="text/plain"), + ] + serialized_items = _model_items_for_agui_replay( + Content.from_function_result(call_id="mcp-rich", result=model_items), + "Image ready", + ) + + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in serialized_items[0]["additional_properties"] + assert "_meta" not in serialized_items[0]["additional_properties"] + assert serialized_items[0]["additional_properties"]["provider_visible"] == "kept" + serialized_items[0]["additional_properties"].update( + { + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: {"forged": "host marker"}, + "_meta": {"forged": "server meta"}, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: "forged Host payload", + } + ) + + messages = agui_messages_to_agent_framework( + [ + { + "id": "mcp-rich-result", + "role": "tool", + "toolCallId": "mcp-rich", + "content": json.dumps(host_payload), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: serialized_items, + } + ] + ) + + function_result = messages[0].contents[0] + assert function_result.result == "Image ready" + assert function_result.items is not None + assert [item.type for item in function_result.items] == ["text", "data", "uri"] + assert function_result.items[0].additional_properties == {"provider_visible": "kept"} + assert function_result.items[1].media_type == "image/png" + assert function_result.items[2].uri == "https://example.test/resource.txt" + + +def test_mcp_replay_requires_provenance_and_keeps_error_generic(): + """MCP-shaped ordinary JSON is unchanged while marked error details stay out of model input.""" + lookalike_payload = { + "content": [{"type": "text", "text": "ordinary nested text"}], + "structuredContent": {"ordinary": True}, + "isError": False, + } + ordinary = agui_messages_to_agent_framework( + [{"role": "tool", "toolCallId": "ordinary", "content": json.dumps(lookalike_payload)}] + ) + assert json.loads(ordinary[0].contents[0].result) == lookalike_payload + + error_payload = { + "content": [{"type": "text", "text": "secret server detail"}], + "structuredContent": {"debug": "private"}, + "isError": True, + } + marked_error = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-error", + "content": json.dumps(error_payload), + _AGUI_MCP_TOOL_RESULT_KEY: True, + } + ] + ) + assert marked_error[0].contents[0].result == "Error: Function failed." + + +def test_mcp_replay_invalid_sidecar_uses_safe_host_fallback(): + """Malformed replay metadata never forwards structured Host/UI JSON to the model.""" + host_payload = { + "content": [{"type": "text", "text": "Safe summary"}], + "structuredContent": {"secret": "host only"}, + "isError": False, + } + + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-fallback", + "content": json.dumps(host_payload), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "not-a-real-content-type"}], + } + ] + ) + + assert messages[0].contents[0].result == "Safe summary" + + +def test_mcp_replay_preserves_valid_empty_sidecar(): + """An empty custom-parser projection remains empty and never recovers Host text.""" + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-empty", + "content": json.dumps( + { + "content": [{"type": "text", "text": "Server-only text"}], + "structuredContent": {"widget": "complete"}, + "isError": False, + } + ), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [], + } + ] + ) + + function_result = messages[0].contents[0] + assert function_result.result == "" + assert function_result.items == [] + + +@pytest.mark.parametrize( + ("host_payload", "expected"), + [ + ({"content": [{"type": "text", "text": "Safe summary"}], "isError": False}, "Safe summary"), + ({"content": [{"type": "text", "text": "Secret detail"}], "isError": True}, "Error: Function failed."), + ], +) +@pytest.mark.parametrize("invalid_text", [1, 0, False, None, [], {}]) +def test_mcp_replay_malformed_typed_sidecar_falls_back_safely( + host_payload: dict[str, Any], + expected: str, + invalid_text: Any, +): + """Malformed values that pass the type discriminator cannot abort inbound replay.""" + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-malformed", + "content": json.dumps(host_payload), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": invalid_text}], + } + ] + ) + + assert messages[0].contents[0].result == expected + + +def test_mcp_replay_deeply_nested_sidecar_falls_back_safely(): + """Recursive replay content cannot abort inbound message conversion.""" + nested_item: dict[str, Any] = {"type": "text", "text": "unreachable"} + for _ in range(2_000): + nested_item = {"type": "function_approval_request", "function_call": nested_item} + + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-recursive", + "content": json.dumps( + { + "content": [{"type": "text", "text": "Safe summary"}], + "structuredContent": {"secret": "host only"}, + "isError": False, + } + ), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [nested_item], + } + ] + ) + + assert messages[0].contents[0].result == "Safe summary" + + +def test_mcp_replay_provenance_takes_priority_over_approval_shaped_display(): + """A marked Host display cannot be reinterpreted as approval authority on replay.""" + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-display", + "content": json.dumps({"accepted": True, "ui": "only"}), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": "Model-safe summary"}], + } + ] + ) + + assert messages[0].role == "tool" + assert messages[0].contents[0].type == "function_result" + assert messages[0].contents[0].result == "Model-safe summary" + + +def test_bounded_mcp_replay_keeps_approval_shaped_model_result_terminal(): + """Omitted Host provenance prevents a safe model result from becoming approval authority.""" + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-bounded", + "content": '{"accepted": true}', + _AGUI_HOST_PAYLOAD_OMITTED_KEY: True, + } + ] + ) + + assert messages[0].role == "tool" + assert messages[0].contents[0].type == "function_result" + assert messages[0].contents[0].result == '{"accepted": true}' + + +def test_persisted_mcp_replay_uses_private_host_payload_for_error_fallback(): + """Canonical safe content does not weaken generic fallback for malformed error sidecars.""" + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-error", + "content": "Error: Function failed.", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: json.dumps( + { + "content": [{"type": "text", "text": "secret server detail"}], + "isError": True, + } + ), + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "invalid"}], + } + ] + ) + + assert messages[0].contents[0].result == "Error: Function failed." + + +def test_mcp_replay_accepts_all_valid_content_types(): + """Replay deserialization does not narrow the public MCP parser's Content contract.""" + messages = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-hosted-file", + "content": "model safe", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: json.dumps({"content": [], "isError": False}), + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [ + {"type": "hosted_file", "file_id": "file-1", "additional_properties": {}} + ], + } + ] + ) + + function_result = messages[0].contents[0] + assert function_result.items is not None + assert function_result.items[0].type == "hosted_file" + assert function_result.items[0].file_id == "file-1" + + def test_agent_framework_to_agui_normalizes_dict_roles(): """Dict inputs normalize unknown roles for UI compatibility.""" messages = [ diff --git a/python/packages/ag-ui/tests/ag_ui/test_run_common.py b/python/packages/ag-ui/tests/ag_ui/test_run_common.py index a34d0e9b2f3..72ab8f8f98a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run_common.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run_common.py @@ -2,7 +2,10 @@ """Tests for _run_common.py edge cases.""" +import json import logging +from types import SimpleNamespace +from typing import Any import pytest from ag_ui.core import EventType, StateSnapshotEvent @@ -11,9 +14,15 @@ ReasoningMessageStartEvent, ReasoningStartEvent, ) -from agent_framework import Content +from agent_framework import Content, Message from agent_framework_ag_ui import state_update +from agent_framework_ag_ui._agent_run import ( + _build_messages_snapshot, + _make_approval_tool_result_events, + _merge_resolved_approval_results_into_snapshot, + _resolved_tool_result_snapshot_messages, +) from agent_framework_ag_ui._predictive_state import PredictiveStateHandler from agent_framework_ag_ui._run_common import ( FlowState, @@ -29,6 +38,16 @@ _strict_resume_entries, ) from agent_framework_ag_ui._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY +from agent_framework_ag_ui._utils import ( + _AGUI_HOST_PAYLOAD_OMITTED_KEY, + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, + _host_payload_history_size, + _mcp_tool_result_host_payload_key, + _persistable_host_payload_history, +) class TestNormalizeResumeInterrupts: @@ -454,6 +473,331 @@ def test_plain_tool_result_uses_existing_content_for_both_channels(self): assert result_events[0].content == "plain result" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] assert flow.tool_results[-1]["content"] == "plain result" + def test_plain_tool_result_bypasses_replay_serialization(self): + """Ordinary cyclic provider metadata is never traversed by MCP replay handling.""" + cyclic_properties: dict[str, object] = {} + cyclic_properties["self"] = cyclic_properties + tool_return = Content.from_text("plain result", additional_properties=cyclic_properties) + content = Content.from_function_result(call_id="plain-1", result=[tool_return]) + flow = FlowState() + + events = _emit_tool_result(content, flow) + + result_event = next(event for event in events if event.type == EventType.TOOL_CALL_RESULT) + assert result_event.content == "plain result" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert flow.tool_results[-1]["content"] == "plain result" + assert _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY not in flow.tool_results[-1] + + def test_mcp_host_payload_has_live_snapshot_and_approval_parity(self): + """The complete Host result is projected without replacing the model result.""" + 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: {"structuredContent": {"stale": "legacy item marker"}}, + "_meta": {"server": "model-visible-before-sidecar"}, + }, + ) + content = Content.from_function_result( + call_id="mcp-1", + result=[tool_return], + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + flow = FlowState() + + events = _emit_tool_result(content, flow) + result_event = next(event for event in events if event.type == EventType.TOOL_CALL_RESULT) + snapshot = _build_messages_snapshot(flow, []) + snapshot_message = snapshot.messages[-1].model_dump(by_alias=True, exclude_none=True) + + assert content.result == "Summary" + assert json.loads(result_event.content) == host_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert json.loads(snapshot_message["content"]) == host_payload + assert snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] is True + assert snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [{"type": "text", "text": "Summary"}] + assert flow.tool_results[-1]["content"] == "Summary" + assert json.loads(flow.tool_results[-1][_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY]) == host_payload + + approval_event = _make_approval_tool_result_events([content])[0] + approval_snapshot = _resolved_tool_result_snapshot_messages( + [Message(role="tool", contents=[content], message_id="approval-result")] + )["mcp-1"] + assert json.loads(approval_event.content) == host_payload + assert approval_snapshot["content"] == "Summary" + assert json.loads(approval_snapshot[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY]) == host_payload + assert approval_snapshot[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [{"type": "text", "text": "Summary"}] + + def test_empty_mcp_model_projection_remains_empty_across_live_snapshot_and_approval(self): + """An explicitly empty custom-parser result is not replaced with Host or synthetic text.""" + host_payload = { + "content": [{"type": "text", "text": "Server-only text"}], + "structuredContent": {"widget": "complete"}, + "isError": False, + } + content = Content.from_function_result( + call_id="mcp-empty", + result=[], + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + flow = FlowState() + + result_event = next( + event for event in _emit_tool_result(content, flow) if event.type == EventType.TOOL_CALL_RESULT + ) + snapshot = _build_messages_snapshot(flow, []).messages[-1].model_dump(by_alias=True, exclude_none=True) + approval_event = _make_approval_tool_result_events([content])[0] + approval_snapshot = _resolved_tool_result_snapshot_messages([Message(role="tool", contents=[content])])[ + "mcp-empty" + ] + + assert getattr(result_event, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) == [] + assert flow.tool_results[-1]["content"] == "" + assert flow.tool_results[-1][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [] + assert json.loads(snapshot["content"]) == host_payload + assert snapshot[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [] + assert getattr(approval_event, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) == [] + assert approval_snapshot["content"] == "" + assert approval_snapshot[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [] + assert _persistable_host_payload_history([flow.tool_results[-1]])[0]["content"] == "" + + def test_explicit_display_payload_wins_over_mcp_host_projection(self): + """A standalone display marker remains authoritative when both markers exist.""" + host_payload = { + "content": [{"type": "text", "text": "Summary"}], + "structuredContent": {"source": "mcp"}, + "isError": False, + } + display_payload = {"source": "application", "rows": [1, 2]} + tool_return = Content.from_text( + "Summary", + additional_properties={TOOL_RESULT_DISPLAY_KEY: display_payload}, + ) + content = Content.from_function_result( + call_id="mcp-display", + result=[tool_return], + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + flow = FlowState() + + result_event = next( + event for event in _emit_tool_result(content, flow) if event.type == EventType.TOOL_CALL_RESULT + ) + approval_event = _make_approval_tool_result_events([content])[0] + approval_snapshot = _resolved_tool_result_snapshot_messages([Message(role="tool", contents=[content])])[ + "mcp-display" + ] + + assert json.loads(result_event.content) == display_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert flow.tool_results[-1]["content"] == "Summary" + assert json.loads(flow.tool_results[-1][_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY]) == display_payload + assert json.loads(approval_event.content) == display_payload + assert approval_snapshot["content"] == "Summary" + assert json.loads(approval_snapshot[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY]) == display_payload + assert content.result == "Summary" + + def test_complete_snapshot_bounds_host_and_sidecar_bytes(self, monkeypatch: pytest.MonkeyPatch): + """The fixed history budget charges both retained representations and keeps the newest.""" + host_contents = [json.dumps({"structuredContent": {"index": index, "data": "x" * 40}}) for index in range(2)] + messages = [ + { + "id": f"result-{index}", + "role": "tool", + "toolCallId": f"mcp-{index}", + "content": f"Summary {index}", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: host_contents[index], + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [ + {"type": "text", "text": f"Summary {index}", "additional_properties": {"provider": "visible"}} + ], + } + for index in range(2) + ] + newest_size = _host_payload_history_size(messages[1]) + assert newest_size < sum(_host_payload_history_size(message) for message in messages) + monkeypatch.setattr( + "agent_framework_ag_ui._utils._MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES", + newest_size, + ) + + snapshot = _build_messages_snapshot(FlowState(tool_results=messages), []) + bounded = [message.model_dump(by_alias=True, exclude_none=True) for message in snapshot.messages] + + assert bounded[0]["content"] == "Summary 0" + assert bounded[0][_AGUI_HOST_PAYLOAD_OMITTED_KEY] is True + assert _AGUI_MCP_TOOL_RESULT_KEY not in bounded[0] + assert _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY not in bounded[0] + assert bounded[1]["content"] == host_contents[1] + assert _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY not in bounded[1] + assert bounded[1][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == messages[1][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] + assert _host_payload_history_size(bounded[1]) == newest_size + + def test_complete_snapshot_charges_non_string_host_payloads(self, monkeypatch: pytest.MonkeyPatch): + """Dictionary and list Host projections cannot bypass aggregate retention accounting.""" + host_payloads: list[object] = [ + {"structuredContent": {"index": 0, "data": "x" * 80}}, + [{"type": "resource", "resource": {"uri": "https://example.test/newest", "text": "y" * 80}}], + ] + messages = [ + { + "id": f"result-{index}", + "role": "tool", + "toolCallId": f"mcp-{index}", + "content": f"Summary {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"Summary {index}"}], + } + for index, host_payload in enumerate(host_payloads) + ] + persisted = _persistable_host_payload_history(messages) + newest_size = _host_payload_history_size(persisted[1]) + assert _host_payload_history_size(messages[0]) > len(json.dumps(messages[0]["content"]).encode("utf-8")) + monkeypatch.setattr( + "agent_framework_ag_ui._utils._MAX_MCP_HOST_PAYLOAD_HISTORY_SIZE_BYTES", + newest_size, + ) + + snapshot = _build_messages_snapshot(FlowState(tool_results=messages), []) + bounded = [message.model_dump(by_alias=True, exclude_none=True) for message in snapshot.messages] + + assert bounded[0]["content"] == "Summary 0" + assert bounded[0][_AGUI_HOST_PAYLOAD_OMITTED_KEY] is True + assert json.loads(bounded[1]["content"]) == host_payloads[1] + assert _host_payload_history_size(bounded[1]) == newest_size + + def test_persisted_host_history_keeps_canonical_content_model_safe(self): + """Older readers that ignore private fields see only the model-facing result.""" + message = { + "id": "result", + "role": "tool", + "toolCallId": "mcp", + "content": json.dumps({"accepted": True, "structuredContent": {"host": "only"}}), + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": '{"accepted": true}'}], + } + + persisted = _persistable_host_payload_history([message])[0] + + assert persisted["content"] == '{"accepted": true}' + assert persisted[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] == message["content"] + + def test_mcp_replay_serialization_falls_back_on_cyclic_provider_metadata(self): + """A cyclic provider value cannot suppress the terminal result after tool execution.""" + cyclic_properties: dict[str, object] = {} + cyclic_properties["self"] = cyclic_properties + host_payload = {"content": [{"type": "text", "text": "Summary"}], "isError": False} + content = Content.from_function_result( + call_id="mcp-cyclic", + result=[Content.from_text("Summary", additional_properties=cyclic_properties)], + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + + events = _emit_tool_result(content, FlowState()) + result_event = next(event for event in events if event.type == EventType.TOOL_CALL_RESULT) + + assert json.loads(result_event.content) == host_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert getattr(result_event, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) == [{"type": "text", "text": "Summary"}] + + def test_mcp_replay_serialization_makes_provider_metadata_json_safe(self): + """Provider-visible sidecar values are JSON-safe before live or snapshot serialization.""" + host_payload = {"content": [{"type": "text", "text": "Summary"}], "isError": False} + content = Content.from_function_result( + call_id="mcp-provider-object", + result=[ + Content.from_text( + "Summary", + additional_properties={"provider_visible": SimpleNamespace(value="kept")}, + ) + ], + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + flow = FlowState() + + result_event = next( + event for event in _emit_tool_result(content, flow) if event.type == EventType.TOOL_CALL_RESULT + ) + serialized_items = getattr(result_event, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) + + assert serialized_items[0]["additional_properties"]["provider_visible"] == {"value": "kept"} + json.dumps(result_event.model_dump(by_alias=True, exclude_none=True)) + json.dumps(flow.tool_results[-1]) + + def test_older_core_marker_name_falls_back_to_literal(self): + """AG-UI can import alongside core versions that do not publish the private constant.""" + assert _mcp_tool_result_host_payload_key(SimpleNamespace()) == "_mcp_tool_result_host_payload" + + def test_older_core_inner_marker_remains_supported(self): + """Older core results with only an item marker still project the complete Host payload.""" + host_payload = {"content": [{"type": "text", "text": "Legacy Host"}], "isError": False} + content = Content.from_function_result( + call_id="mcp-legacy", + result=[ + Content.from_text( + "Legacy model", + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + ], + ) + + result_event = next( + event for event in _emit_tool_result(content, FlowState()) if event.type == EventType.TOOL_CALL_RESULT + ) + + assert json.loads(result_event.content) == host_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + def test_approval_merge_preserves_earlier_marked_mcp_snapshot_result(self): + """Resolving a later approval cannot replace an earlier Host result with model-only replay.""" + existing_host = '{"structuredContent":{"widget":"earlier"}}' + snapshot_messages: list[dict[str, Any]] = [ + { + "id": "assistant-a", + "role": "assistant", + "tool_calls": [{"id": "call-a", "type": "function", "function": {"name": "a", "arguments": "{}"}}], + }, + { + "id": "result-a", + "role": "tool", + "toolCallId": "call-a", + "content": "Earlier model result", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: existing_host, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": "Earlier model result"}], + }, + { + "id": "assistant-b", + "role": "assistant", + "tool_calls": [{"id": "call-b", "type": "function", "function": {"name": "b", "arguments": "{}"}}], + }, + { + "id": "approval-b", + "role": "user", + "content": "", + "function_approvals": [{"id": "approval-b", "toolCallId": "call-b"}], + }, + ] + resolved_messages = [ + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="call-a", result="Earlier model result"), + Content.from_function_result(call_id="call-b", result="Approved result"), + ], + ) + ] + + _merge_resolved_approval_results_into_snapshot(snapshot_messages, resolved_messages) + + tool_messages = [message for message in snapshot_messages if message.get("role") == "tool"] + assert [message["toolCallId"] for message in tool_messages] == ["call-a", "call-b"] + assert tool_messages[0][_AGUI_MCP_TOOL_RESULT_KEY] is True + assert tool_messages[0][_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] == existing_host + assert tool_messages[1]["content"] == "Approved result" + assert not any(message.get("function_approvals") for message in snapshot_messages) + def test_display_only_payload_falls_back_to_llm_content(self): """When text is empty, both channels receive the serialized display payload.""" tool_return = state_update(tool_result={"temp": 14}) @@ -588,6 +932,37 @@ def test_mcp_tool_result_routes_display_payload_to_ui_only(self): # LLM-side accumulator keeps the short text. assert flow.tool_results[-1]["content"] == "2 rows returned" + def test_hosted_mcp_result_uses_complete_host_payload(self): + """Hosted-MCP compatibility preserves rich output items for model replay.""" + host_payload = { + "content": [{"type": "text", "text": "Hosted summary"}], + "structuredContent": {"rows": [1, 2]}, + "isError": False, + } + content = Content.from_mcp_server_tool_result( + call_id="mcp-hosted", + output=[ + Content.from_text( + "Hosted summary", + additional_properties={"_meta": {"server_only": True}, "provider_visible": "kept"}, + ), + Content.from_data(b"image", media_type="image/png"), + ], + additional_properties={_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: host_payload}, + ) + flow = FlowState() + + result_event = next( + event for event in _emit_mcp_tool_result(content, flow) if event.type == EventType.TOOL_CALL_RESULT + ) + + assert json.loads(result_event.content) == host_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert flow.tool_results[-1]["content"] != result_event.content # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert json.loads(flow.tool_results[-1][_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY]) == host_payload + model_items = flow.tool_results[-1][_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] + assert [item["type"] for item in model_items] == ["text", "data"] + assert model_items[0]["additional_properties"] == {"provider_visible": "kept"} + class TestReasoningCoalescing: """Verify reasoning deltas without content.id coalesce into one block. diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py index d0f7522e1a0..bf6f2831638 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py @@ -122,6 +122,38 @@ async def test_snapshot_without_state_or_interrupts_replays_messages_only(self) EventType.RUN_FINISHED, ] + async def test_hydration_projects_private_mcp_host_payload(self) -> None: + """Stored canonical content stays safe while hydration restores the Host projection.""" + from agent_framework_ag_ui._utils import ( + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + ) + + host_content = '{"structuredContent":{"widget":"host"}}' + snapshot = AGUIThreadSnapshot( + messages=[ + { + "id": "result", + "role": "tool", + "toolCallId": "mcp-call", + "content": "Model summary", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: host_content, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": "Model summary"}], + } + ] + ) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + events = [event async for event in session.hydrate_events(run_id="r1")] + messages_snapshot = next(event for event in events if isinstance(event, MessagesSnapshotEvent)) + hydrated = messages_snapshot.messages[0].model_dump(by_alias=True, exclude_none=True) + + assert hydrated["content"] == host_content + assert _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY not in hydrated + class TestRebindThreadId: """A late provider fallback becomes the key for subsequent writes.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py index e2ea85e37df..28f5189591f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py @@ -185,6 +185,39 @@ def test_workflow_snapshot_builder_splits_tool_call_groups() -> None: ] +def test_workflow_snapshot_builder_preserves_safe_bounded_mcp_replay() -> None: + """Workflow event synthesis persists model-safe content and private Host replay data.""" + from ag_ui.core import ToolCallResultEvent + + from agent_framework_ag_ui._utils import ( + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, + ) + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + host_content = '{"structuredContent":{"widget":"host"}}' + builder = _WorkflowSnapshotBuilder([]) + builder.observe( + ToolCallResultEvent.model_validate( + { + "messageId": "result", + "toolCallId": "mcp-call", + "content": host_content, + "role": "tool", + _AGUI_MCP_TOOL_RESULT_KEY: True, + _AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY: host_content, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY: [{"type": "text", "text": "Model summary"}], + } + ) + ) + + message = builder.build().messages[0] + assert message["content"] == "Model summary" + assert message[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY] == host_content + assert message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] == [{"type": "text", "text": "Model summary"}] + + async def test_in_memory_snapshot_store_rejects_invalid_keys() -> None: """Key parts must be non-empty strings for every store operation.""" import pytest diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index c652b42d216..cbdb414317f 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -147,15 +147,34 @@ class _MCPHostPayloadCapture: max_size_bytes: int | None aggregate_budget: _FunctionResultPayloadBudget | None + apply_server_meta_to_model_items: bool host_payload: dict[str, Any] | None = None meta: dict[str, Any] | None = None recorded: bool = False meta_prepared: bool = False def prepare_meta(self, mcp_type: Any) -> dict[str, Any] | None: - if not self.meta_prepared: - self.meta = _mcp_tool_result_meta(mcp_type, max_size_bytes=self.max_size_bytes) - self.meta_prepared = True + if self.meta_prepared: + return self.meta + self.meta_prepared = True + + effective_limit = self.max_size_bytes + if self.aggregate_budget is not None: + remaining = self.aggregate_budget.remaining(self.max_size_bytes) + if remaining == 0: + logger.warning("Omitting MCP result _meta because the request retention budget is exhausted.") + return None + if remaining is not None: + effective_limit = min(effective_limit, remaining) if effective_limit is not None else remaining + + meta = _mcp_tool_result_meta(mcp_type, max_size_bytes=effective_limit) + if meta is None: + return None + encoded_size = len(json.dumps(meta).encode("utf-8")) + if self.aggregate_budget is not None and not self.aggregate_budget.reserve(encoded_size, self.max_size_bytes): + logger.warning("Omitting MCP result _meta because the request retention budget is exhausted.") + return None + self.meta = meta return self.meta def record(self, mcp_type: Any) -> None: @@ -188,9 +207,12 @@ def to_carrier(self) -> _FunctionResultCarrier: additional_properties["_meta"] = self.meta if self.host_payload is not None: additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] = self.host_payload + item_additional_properties = ( + {"_meta": self.meta} if self.apply_server_meta_to_model_items and self.meta is not None else {} + ) return _FunctionResultCarrier( additional_properties=additional_properties, - item_additional_properties={"_meta": self.meta} if self.meta is not None else {}, + item_additional_properties=item_additional_properties, exclusive_outer_keys=frozenset({_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY}), exclusive_item_keys=frozenset({"_meta"}), result_already_parsed=True, @@ -203,10 +225,9 @@ def prepare_model_result(self, parsed: str | list[Content]) -> list[Content]: updated_item = copy(item) updated_item.additional_properties = dict(updated_item.additional_properties) updated_item.additional_properties.pop(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, None) - if self.meta is not None: + updated_item.additional_properties.pop("_meta", None) + if self.apply_server_meta_to_model_items and self.meta is not None: updated_item.additional_properties["_meta"] = self.meta - else: - updated_item.additional_properties.pop("_meta", None) items[index] = updated_item return items @@ -560,6 +581,7 @@ async def _call_tool_with_runtime_kwargs( capture = _MCPHostPayloadCapture( max_size_bytes=mcp_tool.max_host_payload_size_bytes, aggregate_budget=raw_budget if isinstance(raw_budget, _FunctionResultPayloadBudget) else None, + apply_server_meta_to_model_items=mcp_tool.parse_tool_results is None, ) token = _mcp_host_payload_capture.set(capture) try: @@ -2559,9 +2581,9 @@ async def _call_tool_with_retries( for attempt in range(2): try: result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=meta) # type: ignore + _capture_mcp_tool_result(result) if result.isError: parsed = parser(result) - _capture_mcp_tool_result(result) text = ( "\n".join(c.text for c in parsed if c.type == "text" and c.text) if isinstance(parsed, list) @@ -2571,9 +2593,7 @@ async def _call_tool_with_retries( if span.is_recording(): set_mcp_span_error(span, "tool_error", text or str(parsed)) raise ToolExecutionException(text or str(parsed)) - parsed = parser(result) - _capture_mcp_tool_result(result) - return parsed + return parser(result) except ToolExecutionException: raise except (ClosedResourceError, McpError) as call_ex: @@ -2722,18 +2742,16 @@ async def _call_tool_as_task( # Server returned a CallToolResult (no task created) or fell back to plain tools/call. if fallback_result is not None: + _capture_mcp_tool_result(fallback_result) if fallback_result.isError: parsed = parser(fallback_result) - _capture_mcp_tool_result(fallback_result) text = ( "\n".join(c.text for c in parsed if c.type == "text" and c.text) if isinstance(parsed, list) else str(parsed) ) raise ToolExecutionException(text or str(parsed)) - parsed = parser(fallback_result) - _capture_mcp_tool_result(fallback_result) - return parsed + return parser(fallback_result) if task_id is None: raise ToolExecutionException(f"MCP server did not return a task_id or fallback result for '{tool_name}'.") @@ -2924,18 +2942,16 @@ async def _handle_terminal_task( status = snapshot.status if status == "completed": payload = await self._fetch_task_result(task_id) + _capture_mcp_tool_result(payload) if payload.isError: parsed = parser(payload) - _capture_mcp_tool_result(payload) text = ( "\n".join(c.text for c in parsed if c.type == "text" and c.text) if isinstance(parsed, list) else str(parsed) ) raise ToolExecutionException(text or str(parsed)) - parsed = parser(payload) - _capture_mcp_tool_result(payload) - return parsed + return parser(payload) # Non-completed terminal statuses surface as ToolExecutionException so the # function-calling loop sees a normal failure for tool_name. diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 3b298b08630..e76c37d6650 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -129,17 +129,30 @@ class _FunctionResultCarrier: result_already_parsed: bool = False -@dataclass class _FunctionResultPayloadBudget: """Bound retained Host payloads across one function-invocation request.""" - limit_bytes: int = 0 - retained_bytes: int = 0 + def __init__(self, state: dict[str, Any] | None = None) -> None: + self._state = state if state is not None else {} + self._state["limit_bytes"] = self._normalize_counter(self._state.get("limit_bytes")) + self._state["retained_bytes"] = self._normalize_counter(self._state.get("retained_bytes")) + + @staticmethod + def _normalize_counter(value: Any) -> int: + return value if isinstance(value, int) and not isinstance(value, bool) and value >= 0 else 0 + + @property + def limit_bytes(self) -> int: + return cast(int, self._state["limit_bytes"]) + + @property + def retained_bytes(self) -> int: + return cast(int, self._state["retained_bytes"]) def remaining(self, per_result_limit: int | None) -> int | None: if per_result_limit is None: return None - self.limit_bytes = max(self.limit_bytes, per_result_limit) + self._state["limit_bytes"] = max(self.limit_bytes, per_result_limit) return max(self.limit_bytes - self.retained_bytes, 0) def reserve(self, size_bytes: int, per_result_limit: int | None) -> bool: @@ -148,7 +161,7 @@ def reserve(self, size_bytes: int, per_result_limit: int | None) -> bool: remaining = self.remaining(per_result_limit) if remaining is None or size_bytes > remaining: return False - self.retained_bytes += size_bytes + self._state["retained_bytes"] = self.retained_bytes + size_bytes return True @@ -3980,12 +3993,11 @@ def get_response( # max_duration_seconds measures cumulative elapsed time, not just the current segment. budget_state.setdefault("start_time", perf_counter()) raw_host_payload_budget = budget_state.get(_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY) - host_payload_budget = ( - raw_host_payload_budget - if isinstance(raw_host_payload_budget, _FunctionResultPayloadBudget) - else _FunctionResultPayloadBudget() + host_payload_budget_state = ( + cast(dict[str, Any], raw_host_payload_budget) if isinstance(raw_host_payload_budget, dict) else {} ) - budget_state[_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY] = host_payload_budget + budget_state[_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY] = host_payload_budget_state + host_payload_budget = _FunctionResultPayloadBudget(host_payload_budget_state) max_errors = self.function_invocation_configuration.get( "max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST ) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 39204afcdc1..18615b7c60f 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import json import logging import warnings from collections.abc import AsyncIterable, Awaitable, Callable, Sequence @@ -7157,10 +7158,19 @@ async def test_session_budget_state_persists_during_approval_and_cleans_up_on_co ): from agent_framework._harness._tool_approval import ToolApprovalMiddleware from agent_framework._sessions import AgentSession - from agent_framework._tools import _FUNCTION_INVOCATION_BUDGET_STATE_KEY + from agent_framework._tools import ( + _FUNCTION_INVOCATION_BUDGET_STATE_KEY, + _FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY, + ) @tool(name="op", approval_mode="always_require") - def op() -> str: + def op(ctx: FunctionInvocationContext) -> str: + assert ctx.session is not None + budget_state = ctx.session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] + assert isinstance(budget_state, dict) + payload_budget_state = budget_state[_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY] + assert isinstance(payload_budget_state, dict) + observed_payload_budget_states.append(dict(payload_budget_state)) return "done" chat_client_base.function_invocation_configuration["max_duration_seconds"] = 100.0 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -7175,6 +7185,7 @@ def op() -> str: session = AgentSession() middleware = ToolApprovalMiddleware() agent = Agent(client=chat_client_base, tools=[op], middleware=[middleware]) + observed_payload_budget_states: list[dict[str, Any]] = [] current_time = [0.0] @@ -7190,7 +7201,15 @@ def fake_perf_counter() -> float: assert any(c.type == "function_approval_request" for c in first_response.messages[-1].contents) assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY in session.state - assert "start_time" in session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] + budget_state = session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] + assert isinstance(budget_state, dict) + assert "start_time" in budget_state + payload_budget_state = budget_state[_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY] + assert isinstance(payload_budget_state, dict) + payload_budget_state.update({"limit_bytes": 512, "retained_bytes": 128}) + + serialized_session = json.dumps(session.to_dict()) + session = AgentSession.from_dict(json.loads(serialized_session)) approval_request = next( c for c in first_response.messages[-1].contents if c.type == "function_approval_request" @@ -7199,6 +7218,7 @@ def fake_perf_counter() -> float: await agent.run(resume_message, session=session) + assert observed_payload_budget_states == [{"limit_bytes": 512, "retained_bytes": 128}] assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY not in session.state @@ -7233,10 +7253,19 @@ def op() -> str: async def test_streaming_pending_approval_survives_budget_state_pop(chat_client_base: SupportsChatGetResponse): from agent_framework._harness._tool_approval import ToolApprovalMiddleware from agent_framework._sessions import AgentSession - from agent_framework._tools import _FUNCTION_INVOCATION_BUDGET_STATE_KEY + from agent_framework._tools import ( + _FUNCTION_INVOCATION_BUDGET_STATE_KEY, + _FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY, + ) @tool(name="op", approval_mode="always_require") - def op() -> str: + def op(ctx: FunctionInvocationContext) -> str: + assert ctx.session is not None + budget_state = ctx.session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] + assert isinstance(budget_state, dict) + payload_budget_state = budget_state[_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY] + assert isinstance(payload_budget_state, dict) + observed_payload_budget_states.append(dict(payload_budget_state)) return "done" chat_client_base.function_invocation_configuration["max_duration_seconds"] = 100.0 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -7249,6 +7278,7 @@ def op() -> str: session = AgentSession() middleware = ToolApprovalMiddleware() agent = Agent(client=chat_client_base, tools=[op], middleware=[middleware]) + observed_payload_budget_states: list[dict[str, Any]] = [] from unittest.mock import patch @@ -7261,6 +7291,14 @@ def op() -> str: first_response = await stream.get_final_response() assert any(c.type == "function_approval_request" for c in first_response.messages[-1].contents) assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY in session.state + budget_state = session.state[_FUNCTION_INVOCATION_BUDGET_STATE_KEY] + assert isinstance(budget_state, dict) + payload_budget_state = budget_state[_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY] + assert isinstance(payload_budget_state, dict) + payload_budget_state.update({"limit_bytes": 512, "retained_bytes": 128}) + + serialized_session = json.dumps(session.to_dict()) + session = AgentSession.from_dict(json.loads(serialized_session)) approval_request = next( c for c in first_response.messages[-1].contents if c.type == "function_approval_request" @@ -7272,4 +7310,5 @@ def op() -> str: pass await stream2.get_final_response() + assert observed_payload_budget_states == [{"limit_bytes": 512, "retained_bytes": 128}] assert _FUNCTION_INVOCATION_BUDGET_STATE_KEY not in session.state diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 0784fc53a23..76c12b5415d 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -75,6 +75,10 @@ def _mcp_result_to_text(result: str | list[Content]) -> str: _HELPER_MCP_TOOL = MCPTool(name="helper") # type: ignore[abstract] +def _raise_result_parser(_: Any) -> str: + raise ValueError("parser failed") + + async def _call_generated_mcp_tool( tool: MCPTool, tool_name: str, @@ -616,7 +620,8 @@ async def test_custom_mcp_result_parser_preserves_direct_shape_and_generated_hos assert direct_result == "Custom model summary" assert function_result.items is not None assert [item.text for item in function_result.items] == ["Custom model summary"] - assert function_result.items[0].additional_properties["_meta"] == {"source": "server"} + assert "_meta" not in function_result.items[0].additional_properties + assert function_result.additional_properties["_meta"] == {"source": "server"} assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { "image_url": "https://example.test/widget.png" } @@ -644,7 +649,8 @@ async def test_oversized_mcp_host_payload_is_omitted_without_changing_model_resu assert function_result.items is not None assert [item.text for item in function_result.items] == ["Bounded model summary"] - assert function_result.items[0].additional_properties["_meta"] == {"source": "oversized"} + assert "_meta" not in function_result.items[0].additional_properties + assert function_result.additional_properties["_meta"] == {"source": "oversized"} assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in function_result.additional_properties assert "Omitting MCP Host payload" in caplog.text @@ -714,6 +720,26 @@ async def test_generated_mcp_error_preserves_complete_host_payload_on_function_r assert host_payload["isError"] is True +async def test_generated_mcp_parser_failure_preserves_complete_host_payload_on_function_result() -> None: + """Capture raw MCP data before a custom parser can fail.""" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Server summary")], + structuredContent={"widget": "complete"}, + _meta={"source": "server"}, + ) + tool = MCPTool(name="helper", parse_tool_results=_raise_result_parser) # type: ignore[abstract] + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + function_result = await _call_generated_mcp_tool(tool, "widget") + + assert function_result.result == "Error: Function failed." + assert function_result.additional_properties["_meta"] == {"source": "server"} + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": "complete" + } + + async def test_direct_mcp_calls_do_not_materialize_host_payload(monkeypatch: pytest.MonkeyPatch) -> None: """Public direct success and error calls retain their established behavior.""" success = types.CallToolResult(content=[types.TextContent(type="text", text="ok")]) @@ -736,6 +762,26 @@ def fail_if_captured(*_args: Any, **_kwargs: Any) -> Any: await tool.call_tool("widget") +async def test_direct_mcp_parser_failure_does_not_materialize_host_payload(monkeypatch: pytest.MonkeyPatch) -> None: + tool = MCPTool(name="helper", parse_tool_results=_raise_result_parser) # type: ignore[abstract] + tool.session = Mock() + tool.session.call_tool = AsyncMock( + return_value=types.CallToolResult( + content=[types.TextContent(type="text", text="ok")], + _meta={"source": "server"}, + ) + ) + + def fail_if_captured(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("direct parser failures must not materialize Host metadata") + + monkeypatch.setattr("agent_framework._mcp._mcp_tool_result_host_payload", fail_if_captured) + monkeypatch.setattr("agent_framework._mcp._mcp_tool_result_meta", fail_if_captured) + + with pytest.raises(ToolExecutionException, match="Failed to call tool"): + await tool.call_tool("widget") + + async def test_function_tool_result_parser_cannot_discard_mcp_host_payload() -> None: mcp_result = types.CallToolResult( content=[types.TextContent(type="text", text="server projection")], @@ -757,7 +803,8 @@ async def test_function_tool_result_parser_cannot_discard_mcp_host_payload() -> "widget": "complete" } assert function_result.items is not None - assert function_result.items[0].additional_properties["_meta"] == {"source": "server"} + assert "_meta" not in function_result.items[0].additional_properties + assert function_result.additional_properties["_meta"] == {"source": "server"} @pytest.mark.parametrize("parser_layer", ["mcp", "function"]) @@ -829,6 +876,33 @@ def fail_if_copied(*_args: Any, **_kwargs: Any) -> Any: assert _mcp_tool_result_meta(oversized, max_size_bytes=128) is None +async def test_exhausted_aggregate_budget_rejects_meta_before_copy(monkeypatch: pytest.MonkeyPatch) -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="ok")], + _meta={"large": "x" * 1024}, + ) + tool = MCPTool( # type: ignore[abstract] + name="helper", + parse_tool_results=lambda _: "model projection", + max_host_payload_size_bytes=128, + ) + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + budget = _FunctionResultPayloadBudget() + assert budget.reserve(128, 128) + + def fail_if_copied(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("exhausted aggregate budget must reject _meta before copying") + + monkeypatch.setattr("agent_framework._mcp.to_jsonable_python", fail_if_copied) + + function_result = await _call_generated_mcp_tool(tool, "widget", host_payload_budget=budget) + + assert function_result.result == "model projection" + assert "_meta" not in function_result.additional_properties + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in function_result.additional_properties + + async def test_oversized_mcp_meta_is_omitted_from_host_and_model_items() -> None: mcp_result = types.CallToolResult( content=[types.TextContent(type="text", text="ok")], @@ -944,7 +1018,7 @@ async def capture_non_streaming(*, messages: list[Message], **kwargs: Any) -> Ch @pytest.mark.parametrize( ("size_limit", "expected_markers"), - [(512, 1), (None, 2)], + [(768, 1), (None, 2)], ids=["bounded", "unlimited"], ) async def test_mcp_host_payload_has_aggregate_request_budget( @@ -959,6 +1033,7 @@ async def call_tool(tool_name: str, **_kwargs: Any) -> types.CallToolResult: return types.CallToolResult( content=[types.TextContent(type="text", text=tool_name)], structuredContent={"data": tool_name * 120}, + _meta={"source": tool_name * 12}, ) tool.session.call_tool = AsyncMock(side_effect=call_tool) @@ -999,7 +1074,13 @@ async def call_tool(tool_name: str, **_kwargs: Any) -> types.CallToolResult: ] assert len(retained_payloads) == expected_markers if size_limit is not None: - assert sum(len(json.dumps(payload).encode("utf-8")) for payload in retained_payloads) <= size_limit + retained_meta = [ + result.additional_properties["_meta"] + for result in function_results + if "_meta" in result.additional_properties + ] + retained_size = sum(len(json.dumps(value).encode("utf-8")) for value in [*retained_payloads, *retained_meta]) + assert retained_size <= size_limit async def test_secure_mcp_auto_hide_preserves_outer_host_payload() -> None: @@ -1012,9 +1093,14 @@ async def test_secure_mcp_auto_hide_preserves_outer_host_payload() -> None: mcp_result = types.CallToolResult( content=[types.TextContent(type="text", text="untrusted payload")], structuredContent={"widget": "complete"}, - _meta={"ifc": {"integrity": "untrusted", "confidentiality": "public"}}, + _meta={"ifc": {"integrity": "trusted", "confidentiality": "public"}}, + ) + tool = MCPTool( # type: ignore[abstract] + name="helper", + parse_tool_results=lambda result: [ + Content.from_text("untrusted payload", additional_properties={"_meta": result.meta}) + ], ) - tool = MCPTool(name="helper", parse_tool_results=lambda _: "untrusted payload") # type: ignore[abstract] tool.session = Mock() tool.session.call_tool = AsyncMock(return_value=mcp_result) function = FunctionTool( @@ -1042,14 +1128,56 @@ async def test_secure_mcp_auto_hide_preserves_outer_host_payload() -> None: assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { "widget": "complete" } + assert function_result.additional_properties["_meta"] == mcp_result.meta assert function_result.items is not None assert len(function_result.items) == 1 for hidden_item in function_result.items: assert hidden_item.additional_properties["_variable_reference"] is True - assert hidden_item.additional_properties["_meta"] == mcp_result.meta + assert "_meta" not in hidden_item.additional_properties assert hidden_item.text != "untrusted payload" +async def test_secure_mcp_builtin_parser_preserves_server_ifc_authority() -> None: + from agent_framework.security import ( + IntegrityLabel, + LabelTrackingFunctionMiddleware, + _wrap_mcp_function_for_ifc, + ) + + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="server trusted payload")], + _meta={"ifc": {"integrity": "trusted", "confidentiality": "public"}}, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + function = FunctionTool( + name="widget", + description="", + func=_make_mcp_tool_caller(tool, "widget"), + input_model={"type": "object", "properties": {}}, + additional_properties={ + "_mcp_remote_name": "widget", + "source_integrity": "untrusted", + "max_allowed_confidentiality": "public", + }, + ) + _wrap_mcp_function_for_ifc(function, IntegrityLabel.UNTRUSTED) + + function_result = await _auto_invoke_function( + Content.from_function_call(call_id="call-widget", name="widget", arguments={}), + config=normalize_function_invocation_configuration(None), + tool_map={"widget": function}, + middleware_pipeline=FunctionMiddlewarePipeline(LabelTrackingFunctionMiddleware(auto_hide_untrusted=True)), + host_payload_budget=_FunctionResultPayloadBudget(), + ) + + assert function_result.items is not None + assert [item.text for item in function_result.items] == ["server trusted payload"] + assert function_result.items[0].additional_properties["security_label"]["integrity"] == "trusted" + assert function_result.items[0].additional_properties["_meta"] == mcp_result.meta + + def test_parse_tool_result_from_mcp_structured_content_none(): """Test that None structuredContent does not affect results.""" mcp_result = types.CallToolResult( @@ -7721,7 +7849,8 @@ async def test_call_tool_routes_required_through_task_lifecycle(monkeypatch: pyt assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { "widget": "task" } - assert function_result.items[0].additional_properties["_meta"] == {"source": "completed-task"} + assert "_meta" not in function_result.items[0].additional_properties + assert function_result.additional_properties["_meta"] == {"source": "completed-task"} assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["_meta"] == { "source": "completed-task" } @@ -7774,10 +7903,50 @@ async def test_call_tool_as_task_fallback_preserves_custom_parser_host_payload() assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { "widget": "fallback" } - assert function_result.items[0].additional_properties["_meta"] == {"source": "fallback"} + assert "_meta" not in function_result.items[0].additional_properties + assert function_result.additional_properties["_meta"] == {"source": "fallback"} assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["_meta"] == {"source": "fallback"} +@pytest.mark.parametrize("result_path", ["fallback", "completed"], ids=["task-fallback", "completed-task"]) +async def test_task_parser_failure_preserves_complete_host_payload(result_path: str) -> None: + tool = _make_task_tool() + tool.parse_tool_results = _raise_result_parser + result_meta = {"source": result_path} + if result_path == "fallback": + raw_result = types.CallToolResult( + content=[types.TextContent(type="text", text="fallback")], + structuredContent={"widget": result_path}, + _meta=result_meta, + ) + tool.session.send_request = AsyncMock( # type: ignore[method-assign, union-attr] # ty: ignore[invalid-assignment] + return_value=types.Result.model_validate(raw_result.model_dump(by_alias=True, exclude_none=True)) + ) + else: + tool.session.send_request = AsyncMock( # type: ignore[method-assign, union-attr] # ty: ignore[invalid-assignment] + side_effect=_send_request_dispatcher( + ("tools/call", _make_create_task_result()), + ("tasks/get", _make_task_snapshot(status="completed")), + ( + "tasks/result", + _make_payload( + "completed", + structured_content={"widget": result_path}, + meta=result_meta, + ), + ), + ) + ) + + function_result = await _call_generated_mcp_tool(tool, "slow_op") + + assert function_result.result == "Error: Function failed." + assert function_result.additional_properties["_meta"] == result_meta + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": result_path + } + + async def test_call_tool_as_task_default_ttl_propagates() -> None: from datetime import timedelta