diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index d7d46058180..e91e881e5c7 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -604,6 +604,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | AG-UI approval-time follow-up | The full grouped user-input pause remains in message history and emits no synthetic `TOOL_CALL_RESULT`. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approval_follow_up_group_remains_in_history_without_live_tool_result` | | AG-UI approval execution failure | A grouped executor failure becomes one deterministic terminal error result for the approved call. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approval_execution_failure_emits_one_terminal_error_result` | | AG-UI no-approval path | Ordinary tool results do not gain an extra approval result event. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_no_approval_path_emits_no_approval_specific_duplicate_result` | +| AG-UI MCP Host payload | Core preserves bounded successful and error MCP `CallToolResult` payloads separately from built-in or custom model-facing results; oversized Host payloads are rejected by a bounded preflight without changing the model projection. AG-UI projects retained Host payloads consistently through ordinary and approval-resolved live events, aggregate-bounded messages snapshots, and the supported Host-history converter without changing generic outbound requests or replaying UI-only data to the model. | `packages/core/tests/core/test_mcp.py::test_parse_tool_result_from_mcp_preserves_complete_host_payload_once`, `test_custom_mcp_result_parser_preserves_host_payload_and_model_projection`, `test_oversized_mcp_host_payload_is_omitted_without_changing_model_result`, `test_mcp_host_payload_size_preflight_matches_json_and_aborts_before_dump`, `test_mcp_error_preserves_complete_host_payload_on_function_result`, `packages/ag-ui/tests/ag_ui/test_run_common.py::TestEmitToolResultWithState::test_plain_tool_result_does_not_serialize_replay_items`, `test_mcp_host_payload_routes_to_live_event_and_snapshot`, `test_mcp_snapshot_replays_rich_model_items_without_host_payload_duplication`, `test_messages_snapshot_bounds_cumulative_mcp_host_payloads`, `packages/ag-ui/tests/ag_ui/test_message_adapters.py::test_agent_framework_to_agui_preserves_mcp_host_payload_after_reload`, `test_host_history_conversion_is_public_and_bounds_cumulative_payloads`, `test_agui_mcp_fallback_requires_provenance_and_hides_error_details` | | AG-UI client-tool request isolation | Client tool declarations are validated before use and remain request-scoped; a rejected collision or earlier successful request cannot change a later request's server-tool execution. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_failed_client_tool_collision_does_not_affect_next_request`, `test_endpoint_client_tools_do_not_persist_into_next_request` | | AG-UI `confirm_changes` snapshot | An accepted synthetic confirmation is replaced only when its original function call has a real result; rejection is cleaned explicitly, and missing accepted results remain inert. | `packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py` | | AG-UI malformed `confirm_changes` metadata | Non-list tool-call metadata and malformed argument JSON are ignored without guessing a target call. | `test_confirm_changes_target_ignores_non_list_tool_calls`, `test_confirm_changes_target_rejects_malformed_arguments_json` | 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 f577cab705b..a162b7482bd 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,14 @@ ) from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts from ._utils import ( + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, _approval_interrupt_id, + _bound_host_payload_history, _function_call_server_label, + _model_items_for_agui_replay, + _stringify_tool_result, + DEFAULT_MAX_HOST_PAYLOAD_HISTORY_SIZE_BYTES, canonical_function_arguments, convert_agui_tools_to_agent_framework, generate_event_id, @@ -687,7 +693,9 @@ 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) events.append( ToolCallResultEvent( message_id=generate_event_id(), @@ -1968,12 +1976,21 @@ 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": _stringify_tool_result(host_payload) if has_host_payload else llm_result, } + if has_host_payload: + snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True + 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 @@ -2067,7 +2084,12 @@ def _build_messages_snapshot( if flow.snapshot_segments: _append_segmented_snapshot_messages(flow, all_messages) - return MessagesSnapshotEvent(messages=all_messages) # type: ignore[arg-type] + return MessagesSnapshotEvent( + messages=_bound_host_payload_history( # type: ignore[arg-type] + all_messages, + max_size_bytes=DEFAULT_MAX_HOST_PAYLOAD_HISTORY_SIZE_BYTES, + ) + ) # type: ignore[arg-type] # Add assistant message with tool calls only (no content) if flow.pending_tool_calls: @@ -2101,7 +2123,12 @@ 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] + return MessagesSnapshotEvent( + messages=_bound_host_payload_history( # type: ignore[arg-type] + all_messages, + max_size_bytes=DEFAULT_MAX_HOST_PAYLOAD_HISTORY_SIZE_BYTES, + ) + ) # type: ignore[arg-type] def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str, Any]]: 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 48e2f45e350..a2517d5970d 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 @@ -15,9 +15,19 @@ Message, ) +from ._state import TOOL_RESULT_DISPLAY_KEY from ._utils import ( + _AGUI_MCP_TOOL_RESULT_KEY, + _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, AGUI_TO_FRAMEWORK_ROLE, + DEFAULT_MAX_HOST_PAYLOAD_HISTORY_SIZE_BYTES, FRAMEWORK_TO_AGUI_ROLE, + _bound_host_payload_history, + _extract_mcp_tool_result_host_payload, + _extract_tool_result_marker_values, + _model_content_from_mcp_host_payload, + _model_items_for_agui_replay, + _stringify_tool_result, get_role_value, normalize_agui_role, safe_json_parse, @@ -845,6 +855,29 @@ def _filter_modified_args( result.append(chat_msg) continue + if msg.get(_AGUI_MCP_TOOL_RESULT_KEY) is True: + serialized_items = msg.get(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) + model_items: list[Content] | None = None + if isinstance(serialized_items, list): + try: + model_items = [ + Content.from_dict(item) + for item in serialized_items + if isinstance(item, dict) and item.get("type") + ] + except (TypeError, ValueError): + model_items = None + if not model_items: + model_items = [Content.from_text(_model_content_from_mcp_host_payload(parsed))] + chat_msg = Message( + role="tool", + contents=[Content.from_function_result(call_id=str(tool_call_id), result=model_items)], + ) + if "id" in msg: + chat_msg.message_id = msg["id"] + result.append(chat_msg) + continue + # Cast result_content to acceptable type for function_result content func_result: str | dict[str, Any] | list[Any] if isinstance(result_content, str): @@ -958,15 +991,13 @@ def _filter_modified_args( return result -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) +def _convert_agent_framework_messages_to_agui( + messages: list[Message] | list[dict[str, Any]], + *, + include_host_payload: bool, +) -> list[dict[str, Any]]: + """Convert Agent Framework messages to AG-UI format.""" - Returns: - List of AG-UI message dictionaries - """ from ._utils import generate_event_id result: list[dict[str, Any]] = [] @@ -1025,14 +1056,24 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An # id (e.g. f"{base_id}-1") risks colliding with a legitimate id elsewhere # in the history, which would let id-keyed clients re-collapse results. for idx, fr in enumerate(function_results): - result.append( - { - "id": msg.message_id if (idx == 0 and msg.message_id) else generate_event_id(), - "role": "tool", - "content": fr.result if fr.result is not None else "", - "toolCallId": fr.call_id, - } - ) + model_result = fr.result if fr.result is not None else "" + tool_message: dict[str, Any] = { + "id": msg.message_id if (idx == 0 and msg.message_id) else generate_event_id(), + "role": "tool", + "content": model_result, + "toolCallId": fr.call_id, + } + has_host_payload, host_payload = _extract_mcp_tool_result_host_payload(fr) + if include_host_payload and has_host_payload: + display_values = _extract_tool_result_marker_values(fr, TOOL_RESULT_DISPLAY_KEY) + tool_message["content"] = _stringify_tool_result( + display_values[-1] if display_values else host_payload + ) + tool_message[_AGUI_MCP_TOOL_RESULT_KEY] = True + tool_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = _model_items_for_agui_replay( + fr, _stringify_tool_result(model_result) + ) + result.append(tool_message) # A mixed message may also carry text / function_call contents alongside # the tool results (e.g. a finalized assistant turn). Emit those as a # separate, distinctly-identified message so they are not lost. @@ -1061,6 +1102,24 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An 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, include_host_payload=False) + + +def agent_framework_messages_to_agui_host_history( + messages: list[Message] | list[dict[str, Any]], + *, + max_host_payload_history_size_bytes: int = DEFAULT_MAX_HOST_PAYLOAD_HISTORY_SIZE_BYTES, +) -> list[dict[str, Any]]: + """Convert Agent Framework messages to bounded AG-UI Host history with replay metadata.""" + converted = _convert_agent_framework_messages_to_agui(messages, include_host_payload=True) + return _bound_host_payload_history( + converted, + max_size_bytes=max_host_payload_history_size_bytes, + ) + + 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 5f73c58d53b..afd0050d046 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,18 @@ 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_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 +707,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,8 +740,12 @@ 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_tool_result_host_payload(content: Content, display_result: Any) -> tuple[bool, Any]: # noqa: ANN401 + """Resolve an MCP Host payload while retaining an explicitly authored display override.""" + 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 _resolve_ui_payload(llm_str: str, display_result: Any) -> str: # noqa: ANN401 @@ -762,6 +761,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 +790,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 +801,19 @@ 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": snapshot_result_content, + } + if snapshot_result is not _UNSET: + # Snapshots can be resubmitted as model history, so retain the original model projection out of band. + snapshot_message[_AGUI_MCP_TOOL_RESULT_KEY] = True + snapshot_message[_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY] = model_items or [ + {"type": "text", "text": result_content} + ] + 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 +858,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 +868,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 +1036,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 +1046,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/_utils.py b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py index 62d4f39c26f..59dad0f4ad2 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py @@ -11,8 +11,20 @@ from typing import Any from agent_framework import AgentResponseUpdate, ChatResponseUpdate, FunctionTool +from agent_framework import _mcp as _core_mcp # pyright: ignore[reportPrivateUsage] from agent_framework._serialization import make_json_safe # pyright: ignore[reportPrivateUsage] +# AG-UI supports older core releases that predate this marker; those versions retain the existing fallback behavior. +_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY = getattr( + _core_mcp, + "_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY", + "_mcp_tool_result_host_payload", +) +_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY = "_agentFrameworkModelContent" +_AGUI_MCP_TOOL_RESULT_KEY = "_agentFrameworkMcpResult" +_AGUI_HOST_PAYLOAD_OMITTED_KEY = "_agentFrameworkHostPayloadOmitted" +DEFAULT_MAX_HOST_PAYLOAD_HISTORY_SIZE_BYTES = 8 * 1024 * 1024 + # Role mapping constants AGUI_TO_FRAMEWORK_ROLE: dict[str, str] = { "user": "user", @@ -55,6 +67,124 @@ 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.""" + values = _extract_tool_result_marker_values(content, _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 an MCP snapshot loses its model-content 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 repeating the MCP Host payload marker.""" + items = getattr(content, "items", None) or () + if not items: + return [{"type": "text", "text": model_result}] + + serialized_items: list[dict[str, Any]] = [] + for item in items: + serialized = item.to_dict() + additional_properties = serialized.get("additional_properties") + if isinstance(additional_properties, dict): + additional_properties.pop(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, None) + if not additional_properties: + serialized.pop("additional_properties", None) + serialized_items.append(serialized) + return serialized_items + + +def _model_text_from_replay_items(message: dict[str, Any]) -> str: + serialized_items = message.get(_AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) + if not isinstance(serialized_items, list): + return "" + return "\n".join( + item["text"] + for item in serialized_items + if isinstance(item, dict) and item.get("type") == "text" and isinstance(item.get("text"), str) + ) + + +def _bound_host_payload_history( + messages: list[dict[str, Any]], + *, + max_size_bytes: int, +) -> list[dict[str, Any]]: + """Retain the newest MCP Host payloads within one aggregate history budget.""" + if max_size_bytes < 0: + raise ValueError("max_size_bytes must be non-negative.") + + 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 + content = message.get("content") + if not isinstance(content, str): + continue + content_size = len(content.encode("utf-8")) + if retained_size + content_size > max_size_bytes: + omit_indices.add(index) + else: + retained_size += content_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[_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/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index 1749cafe724..2ff2d164c21 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 @@ -9,9 +9,11 @@ import pytest from agent_framework import Content, Message +from agent_framework._mcp import _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY 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, @@ -49,6 +51,107 @@ 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(): + """History conversion uses persisted MCP Host data without changing 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: 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 "_agentFrameworkMcpResult" not in outbound[0] + assert json.loads(converted[0]["content"]) == host_payload + assert converted[0]["toolCallId"] == "mcp-1" + + provider_messages = agui_messages_to_agent_framework(converted) + assert provider_messages[0].contents[0].result == "Summary" + + converted[0].pop("_agentFrameworkModelContent") + fallback_provider_messages = agui_messages_to_agent_framework(converted) + assert fallback_provider_messages[0].contents[0].result == "Summary" + + +def test_host_history_conversion_is_public_and_bounds_cumulative_payloads(): + """The supported converter keeps newest Host data within its 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) + converted = package_converter( + messages, + max_host_payload_history_size_bytes=len(unbounded[1]["content"].encode("utf-8")), + ) + + assert namespace_converter is package_converter + assert converted[0]["content"] == "Summary 0" + assert converted[0]["_agentFrameworkHostPayloadOmitted"] is True + assert json.loads(converted[1]["content"])["structuredContent"]["index"] == 1 + + +def test_agui_mcp_fallback_requires_provenance_and_hides_error_details(): + """Only marked MCP history is reconstructed, and errors keep their generic model projection.""" + lookalike_payload = { + "content": [{"type": "text", "text": "ordinary nested text"}], + "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 + + mcp_error = agui_messages_to_agent_framework( + [ + { + "role": "tool", + "toolCallId": "mcp-error", + "content": json.dumps( + { + "content": [{"type": "text", "text": "secret server detail"}], + "isError": True, + } + ), + "_agentFrameworkMcpResult": True, + } + ] + ) + assert mcp_error[0].contents[0].result == "Error: Function failed." + + 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..111e6f68be6 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,6 +2,7 @@ """Tests for _run_common.py edge cases.""" +import json import logging import pytest @@ -11,9 +12,16 @@ ReasoningMessageStartEvent, ReasoningStartEvent, ) -from agent_framework import Content +from agent_framework import Content, Message +from agent_framework._mcp import _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY from agent_framework_ag_ui import state_update +from agent_framework_ag_ui._agent_run import ( + _build_messages_snapshot, + _make_approval_tool_result_events, + _resolved_tool_result_snapshot_messages, +) +from agent_framework_ag_ui._message_adapters import normalize_agui_input_messages from agent_framework_ag_ui._predictive_state import PredictiveStateHandler from agent_framework_ag_ui._run_common import ( FlowState, @@ -454,6 +462,117 @@ 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_does_not_serialize_replay_items(self): + """Ordinary results bypass MCP replay serialization, including cyclic provider metadata.""" + 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 "_agentFrameworkModelContent" not in flow.tool_results[-1] + + def test_mcp_host_payload_routes_to_live_event_and_snapshot(self): + """MCP Host data replaces neither the model result nor either AG-UI Host surface.""" + 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}, + ) + content = Content.from_function_result(call_id="mcp-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) + snapshot = _build_messages_snapshot(flow, []) + + assert content.result == "Summary" + assert json.loads(result_event.content) == host_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert json.loads(flow.tool_results[-1]["content"]) == host_payload + snapshot_content = snapshot.messages[-1].content + assert isinstance(snapshot_content, str) + assert json.loads(snapshot_content) == host_payload + + snapshot_messages = [message.model_dump(by_alias=True, exclude_none=True) for message in snapshot.messages] + provider_messages, _ = normalize_agui_input_messages(snapshot_messages, sanitize_tool_history=False) + assert provider_messages[-1].contents[0].result == "Summary" + + approval_event = _make_approval_tool_result_events([content])[0] + assert json.loads(approval_event.content) == host_payload + approval_snapshot = _resolved_tool_result_snapshot_messages( + [Message(role="tool", contents=[content], message_id="approval-result")] + ) + assert json.loads(approval_snapshot["mcp-1"]["content"]) == host_payload + assert approval_snapshot["mcp-1"]["_agentFrameworkModelContent"] == [{"type": "text", "text": "Summary"}] + + def test_mcp_snapshot_replays_rich_model_items_without_host_payload_duplication(self): + """The replay sidecar retains model media and omits the separate Host payload.""" + 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}, + ), + Content.from_data(b"image", media_type="image/png"), + ] + content = Content.from_function_result(call_id="mcp-rich", result=model_items) + flow = FlowState() + + _emit_tool_result(content, flow) + + sidecar = flow.tool_results[-1]["_agentFrameworkModelContent"] + assert [item["type"] for item in sidecar] == ["text", "data"] + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in sidecar[0].get("additional_properties", {}) + provider_messages, _ = normalize_agui_input_messages(flow.tool_results, sanitize_tool_history=False) + restored_items = provider_messages[0].contents[0].items + assert restored_items is not None + assert [item.type for item in restored_items] == ["text", "data"] + assert restored_items[1].media_type == "image/png" + + def test_messages_snapshot_bounds_cumulative_mcp_host_payloads( + self, + monkeypatch: pytest.MonkeyPatch, + ): + """Complete thread snapshots retain only the newest Host payloads within the aggregate budget.""" + host_contents = [json.dumps({"structuredContent": {"index": index, "data": "x" * 64}}) for index in range(2)] + flow = FlowState( + tool_results=[ + { + "id": f"result-{index}", + "role": "tool", + "toolCallId": f"mcp-{index}", + "content": host_content, + "_agentFrameworkMcpResult": True, + "_agentFrameworkModelContent": [{"type": "text", "text": f"Summary {index}"}], + } + for index, host_content in enumerate(host_contents) + ] + ) + monkeypatch.setattr( + "agent_framework_ag_ui._agent_run.DEFAULT_MAX_HOST_PAYLOAD_HISTORY_SIZE_BYTES", + len(host_contents[1].encode("utf-8")), + ) + + snapshot = _build_messages_snapshot(flow, []) + messages = [message.model_dump(by_alias=True, exclude_none=True) for message in snapshot.messages] + + assert messages[0]["content"] == "Summary 0" + assert messages[0]["_agentFrameworkHostPayloadOmitted"] is True + assert messages[1]["content"] == host_contents[1] + 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}) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 3d1d5717e37..f42dd13bb92 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -14,8 +14,9 @@ from abc import abstractmethod from collections.abc import Callable, Collection, Coroutine, Mapping, Sequence from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore +from copy import copy from dataclasses import dataclass -from datetime import timedelta +from datetime import date, datetime, timedelta from functools import partial from inspect import isawaitable from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict, cast @@ -29,6 +30,7 @@ _warn_on_feature_use, # pyright: ignore[reportPrivateUsage] experimental, ) +from ._serialization import make_json_safe from ._telemetry import FeatureIndex, mark_feature_used from ._tools import FunctionTool from ._types import ( @@ -90,6 +92,7 @@ class MCPSpecificApproval(TypedDict, total=False): _MCP_REMOTE_NAME_KEY = "_mcp_remote_name" _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name" +_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY = "_mcp_tool_result_host_payload" _MCP_PROGRESSIVE_LIST_TOOL_NAME = "list_mcp_tools" _MCP_PROGRESSIVE_LOAD_TOOL_NAME = "load_tool" _MCP_PROGRESSIVE_UNLOAD_TOOL_NAME = "unload_tool" @@ -124,8 +127,192 @@ class MCPSpecificApproval(TypedDict, total=False): "_meta", }) _mcp_call_headers: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("_mcp_call_headers") +_preserve_mcp_host_payload: contextvars.ContextVar[bool] = contextvars.ContextVar( + "_preserve_mcp_host_payload", + default=False, +) MCP_DEFAULT_TIMEOUT = 30 MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5 +_DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES = 1024 * 1024 + + +class _EncodedSizeBudget: + """Track JSON bytes and abort as soon as an untrusted value exceeds its budget.""" + + def __init__(self, limit: int) -> None: + self.remaining = limit + + def consume(self, size: int) -> None: + self.remaining -= size + if self.remaining < 0: + raise OverflowError + + +def _consume_json_string(value: str, budget: _EncodedSizeBudget) -> None: + budget.consume(2) + for char in value: + codepoint = ord(char) + if char in {'"', "\\"} or char in {"\b", "\f", "\n", "\r", "\t"}: + budget.consume(2) + elif codepoint < 0x20 or codepoint > 0xFFFF: + budget.consume(6 if codepoint <= 0xFFFF else 12) + elif codepoint > 0x7F: + budget.consume(6) + else: + budget.consume(1) + + +def _consume_json_size(value: Any, budget: _EncodedSizeBudget) -> None: + if value is None: + budget.consume(4) + return + if value is True: + budget.consume(4) + return + if value is False: + budget.consume(5) + return + if isinstance(value, str): + _consume_json_string(value, budget) + return + if isinstance(value, (bytes, bytearray)): + budget.consume(2 + 4 * ((len(value) + 2) // 3)) + return + if isinstance(value, (int, float)): + budget.consume(len(json.dumps(value))) + return + if isinstance(value, (datetime, date)): + _consume_json_string(value.isoformat(), budget) + return + if isinstance(value, Mapping): + budget.consume(2) + value_mapping = cast(Mapping[Any, Any], value) + for index, (key, item) in enumerate(value_mapping.items()): + if index: + budget.consume(2) + _consume_json_string(str(key), budget) + budget.consume(2) + _consume_json_size(item, budget) + return + if isinstance(value, Sequence): + budget.consume(2) + value_sequence = cast(Sequence[Any], value) + for index, item in enumerate(value_sequence): + if index: + budget.consume(2) + _consume_json_size(item, budget) + return + + model_fields_raw = getattr(value.__class__, "model_fields", None) + if isinstance(model_fields_raw, Mapping): + model_fields = cast(Mapping[str, Any], model_fields_raw) + budget.consume(2) + field_count = 0 + for name, field_info in model_fields.items(): + item = getattr(value, name, None) + if item is None or getattr(field_info, "exclude", False): + continue + if field_count: + budget.consume(2) + alias = getattr(field_info, "serialization_alias", None) or getattr(field_info, "alias", None) or name + _consume_json_string(str(alias), budget) + budget.consume(2) + _consume_json_size(item, budget) + field_count += 1 + model_extra_raw = getattr(value, "model_extra", None) + if isinstance(model_extra_raw, Mapping): + model_extra = cast(Mapping[Any, Any], model_extra_raw) + for key, item in model_extra.items(): + if item is None: + continue + if field_count: + budget.consume(2) + _consume_json_string(str(key), budget) + budget.consume(2) + _consume_json_size(item, budget) + field_count += 1 + return + + _consume_json_string(str(value), budget) + + +def _json_size_exceeds(value: Any, limit: int) -> bool: + try: + _consume_json_size(value, _EncodedSizeBudget(limit)) + except OverflowError: + return True + return False + + +def _mcp_tool_result_host_payload( + mcp_type: Any, + *, + max_size_bytes: int | None, +) -> dict[str, Any] | None: + """Return a JSON-safe complete MCP result when the value supports model serialization.""" + model_dump = getattr(mcp_type, "model_dump", None) + if not callable(model_dump): + return None + if max_size_bytes is not None and _json_size_exceeds(mcp_type, max_size_bytes): + logger.warning( + "Omitting MCP Host payload because its encoded size exceeds max_host_payload_size_bytes=%d.", + max_size_bytes, + ) + return None + try: + dumped = model_dump(by_alias=True, exclude_none=True, mode="json", fallback=str) + except ValueError: + dumped = model_dump(by_alias=True, exclude_none=True) + if not isinstance(dumped, Mapping): + return None + host_payload = cast(dict[str, Any], make_json_safe(dict(cast(Mapping[str, Any], dumped)))) + if max_size_bytes is not None and len(json.dumps(host_payload).encode("utf-8")) > max_size_bytes: + logger.warning( + "Omitting MCP Host payload because its encoded size exceeds max_host_payload_size_bytes=%d.", + max_size_bytes, + ) + return None + return host_payload + + +def _with_mcp_tool_result_host_payload( + parsed: str | list[Content], + mcp_type: Any, + *, + max_size_bytes: int | None, +) -> list[Content]: + """Attach the complete MCP result once without changing the parser's model projection.""" + items = [Content.from_text(parsed)] if isinstance(parsed, str) else list(parsed) + if not items: + items = [Content.from_text("[]")] + + host_payload = _mcp_tool_result_host_payload(mcp_type, max_size_bytes=max_size_bytes) + raw_meta = getattr(mcp_type, "meta", None) + meta = dict(cast(Mapping[str, Any], raw_meta)) if isinstance(raw_meta, Mapping) else None + for index, item in enumerate(items): + if not (isinstance(meta, Mapping) or (index == 0 and host_payload is not None)): + continue + updated_item = copy(item) + updated_item.additional_properties = dict(updated_item.additional_properties) + if isinstance(meta, Mapping): + updated_item.additional_properties["_meta"] = dict(cast(Mapping[str, Any], meta)) + if index == 0 and host_payload is not None: + updated_item.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] = host_payload + items[index] = updated_item + return items + + +class _MCPToolResultException(ToolExecutionException): + """Carry an MCP Host payload through generic function error conversion.""" + + def __init__(self, message: str, host_payload: dict[str, Any] | None) -> None: + super().__init__(message) + self._function_result_additional_properties: dict[str, Any] = {} + if host_payload is not None: + self._function_result_additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] = host_payload + if isinstance(meta := host_payload.get("_meta"), Mapping): + self._function_result_additional_properties["_meta"] = dict(cast(Mapping[str, Any], meta)) + # Default safety limits applied to server-initiated MCP sampling requests # (``sampling/createMessage``). MCP servers are untrusted third parties, so the @@ -263,7 +450,11 @@ async def _call_tool_with_runtime_kwargs( call_kwargs["_meta"] = trusted_meta else: call_kwargs.pop("_meta", None) - return await mcp_tool.call_tool(remote_tool_name, **call_kwargs) + token = _preserve_mcp_host_payload.set(True) + try: + return await mcp_tool.call_tool(remote_tool_name, **call_kwargs) + finally: + _preserve_mcp_host_payload.reset(token) return _call_tool_with_runtime_kwargs @@ -456,6 +647,7 @@ def __init__( sampling_approval_callback: SamplingApprovalCallback | None = None, sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS, sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, additional_properties: dict[str, Any] | None = None, task_options: MCPTaskOptions | None = None, additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, @@ -513,6 +705,9 @@ def __init__( sampling_max_requests: Maximum number of sampling requests allowed per session connection; further requests are rejected. The counter resets on reconnect. Set to ``None`` to disable the limit. Defaults to ``_DEFAULT_SAMPLING_MAX_REQUESTS``. + max_host_payload_size_bytes: Maximum encoded size of the complete MCP result retained + for Host transports. Oversized payloads are omitted from the Host channel while + the parsed model result is preserved. Set to ``None`` to disable the limit. additional_properties: Additional properties for the tool. task_options: Options controlling how long-running MCP tasks are driven for tools that advertise ``execution.taskSupport == "required"``. When ``None``, @@ -536,6 +731,8 @@ def __init__( object_name="MCP progressive disclosure", category=ExperimentalWarning, ) + if max_host_payload_size_bytes is not None and max_host_payload_size_bytes <= 0: + raise ValueError("max_host_payload_size_bytes must be positive or None.") self.name = name self.description = description or "" self.approval_mode = approval_mode @@ -546,6 +743,7 @@ def __init__( self.parse_tool_results = parse_tool_results self.load_prompts_flag = load_prompts self.parse_prompt_results = parse_prompt_results + self.max_host_payload_size_bytes = max_host_payload_size_bytes # Defer constructing the default MCPTaskOptions so the experimental warning # only fires when LRO is actually engaged (lazy-resolved by _effective_task_options). self._task_options_explicit: MCPTaskOptions | None = task_options @@ -666,13 +864,15 @@ def _parse_tool_result_from_mcp( to derive per-item security labels. The sentinel is intentionally generic so any MCP server's ``_meta`` keys (current or future) can be interpreted by higher-level code. + + The complete MCP result is also preserved under a core-owned + ``additional_properties`` marker. Host transports can use that payload + without changing the content selected for the model. """ from mcp import types raw_meta = mcp_type.meta meta: dict[str, Any] | None = dict(raw_meta) if isinstance(raw_meta, Mapping) else None - # Stamp the server ``_meta`` payload directly via additional_properties on - # each newly constructed Content; empty when the server provided no meta. additional_kwargs: dict[str, Any] = {"additional_properties": {"_meta": meta}} if meta else {} result: list[Content] = [] @@ -717,7 +917,7 @@ def _parse_tool_result_from_mcp( result.append(Content.from_text(str(item), **additional_kwargs)) if mcp_type.structuredContent is not None: - result.append(Content.from_text(json.dumps(mcp_type.structuredContent, default=str))) + result.append(Content.from_text(json.dumps(mcp_type.structuredContent, default=str), **additional_kwargs)) if not result: result.append(Content.from_text("null", **additional_kwargs)) @@ -2070,6 +2270,19 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: ToolExecutionException: If the MCP server is not connected, tools are not loaded, or the tool call fails. """ + return await self._call_tool( + tool_name, + kwargs, + preserve_host_payload=_preserve_mcp_host_payload.get(), + ) + + async def _call_tool( + self, + tool_name: str, + kwargs: dict[str, Any], + *, + preserve_host_payload: bool, + ) -> str | list[Content]: if not self.load_tools_flag: raise ToolExecutionException( "Tools are not loaded for this server, please set load_tools=True in the constructor." @@ -2078,7 +2291,7 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: # Tools advertising taskSupport == "required" cannot complete via plain tools/call; # route through the long-running task lifecycle transparently. if self._tool_task_support_by_name.get(tool_name) == "required": - return await self.call_tool_as_task(tool_name, **kwargs) + return await self._call_tool_as_task(tool_name, kwargs, preserve_host_payload=preserve_host_payload) filtered_kwargs, meta = self._prepare_call_kwargs(tool_name, kwargs) @@ -2091,7 +2304,14 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION, }) with create_mcp_client_span("tools/call", target=tool_name, attributes=mcp_span_attrs) as span: - return await self._call_tool_with_retries(tool_name, filtered_kwargs, meta, parser, span) + return await self._call_tool_with_retries( + tool_name, + filtered_kwargs, + meta, + parser, + span, + preserve_host_payload=preserve_host_payload, + ) async def _call_tool_with_retries( self, @@ -2100,6 +2320,8 @@ async def _call_tool_with_retries( meta: dict[str, Any] | None, parser: Callable[..., str | list[Content]], span: otel_trace.Span, + *, + preserve_host_payload: bool, ) -> str | list[Content]: """Execute the MCP tools/call RPC with retry logic.""" from anyio import ClosedResourceError @@ -2118,8 +2340,23 @@ async def _call_tool_with_retries( # Per OTel MCP semconv: set error.type="tool_error" for isError results if span.is_recording(): set_mcp_span_error(span, "tool_error", text or str(parsed)) - raise ToolExecutionException(text or str(parsed)) - return parser(result) + raise _MCPToolResultException( + text or str(parsed), + _mcp_tool_result_host_payload( + result, + max_size_bytes=self.max_host_payload_size_bytes, + ), + ) + parsed = parser(result) + return ( + _with_mcp_tool_result_host_payload( + parsed, + result, + max_size_bytes=self.max_host_payload_size_bytes, + ) + if preserve_host_payload + else parsed + ) except ToolExecutionException: raise except (ClosedResourceError, McpError) as call_ex: @@ -2229,6 +2466,15 @@ async def call_tool_as_task(self, tool_name: str, **kwargs: Any) -> str | list[C A list of Content items (or a string when a custom ``parse_tool_results`` callback is configured). """ + return await self._call_tool_as_task(tool_name, kwargs, preserve_host_payload=False) + + async def _call_tool_as_task( + self, + tool_name: str, + kwargs: dict[str, Any], + *, + preserve_host_payload: bool, + ) -> str | list[Content]: from anyio import ClosedResourceError from mcp.shared.exceptions import McpError @@ -2268,8 +2514,23 @@ async def call_tool_as_task(self, tool_name: str, **kwargs: Any) -> str | list[C if isinstance(parsed, list) else str(parsed) ) - raise ToolExecutionException(text or str(parsed)) - return parser(fallback_result) + raise _MCPToolResultException( + text or str(parsed), + _mcp_tool_result_host_payload( + fallback_result, + max_size_bytes=self.max_host_payload_size_bytes, + ), + ) + parsed = parser(fallback_result) + return ( + _with_mcp_tool_result_host_payload( + parsed, + fallback_result, + max_size_bytes=self.max_host_payload_size_bytes, + ) + if preserve_host_payload + else parsed + ) if task_id is None: raise ToolExecutionException(f"MCP server did not return a task_id or fallback result for '{tool_name}'.") @@ -2281,7 +2542,13 @@ async def call_tool_as_task(self, tool_name: str, **kwargs: Any) -> str | list[C async def _await_task_completion() -> str | list[Content]: terminal = await self._poll_task_until_terminal(task_id) - return await self._handle_terminal_task(tool_name, task_id, terminal, parser) + return await self._handle_terminal_task( + tool_name, + task_id, + terminal, + parser, + preserve_host_payload=preserve_host_payload, + ) try: if max_wait_s is not None: @@ -2361,7 +2628,6 @@ async def _call_tool_as_task_create( # Inspect the raw payload: a CreateTaskResult carries `task.taskId`; # a legacy CallToolResult carries `content` and/or `isError`. raw: dict[str, Any] = lenient.model_dump(by_alias=True, exclude_none=True) - raw.pop("_meta", None) task_field = raw.get("task") if isinstance(task_field, dict): @@ -2451,6 +2717,8 @@ async def _handle_terminal_task( task_id: str, snapshot: types.GetTaskResult, parser: Callable[[types.CallToolResult], str | list[Content]], + *, + preserve_host_payload: bool, ) -> str | list[Content]: """Map a terminal task snapshot to either a parsed result or an exception.""" status = snapshot.status @@ -2463,8 +2731,23 @@ async def _handle_terminal_task( if isinstance(parsed, list) else str(parsed) ) - raise ToolExecutionException(text or str(parsed)) - return parser(payload) + raise _MCPToolResultException( + text or str(parsed), + _mcp_tool_result_host_payload( + payload, + max_size_bytes=self.max_host_payload_size_bytes, + ), + ) + parsed = parser(payload) + return ( + _with_mcp_tool_result_host_payload( + parsed, + payload, + max_size_bytes=self.max_host_payload_size_bytes, + ) + if preserve_host_payload + else parsed + ) # Non-completed terminal statuses surface as ToolExecutionException so the # function-calling loop sees a normal failure for tool_name. @@ -2496,7 +2779,6 @@ async def _fetch_task_result(self, task_id: str) -> types.CallToolResult: # GetTaskPayloadResult carries the tool result via extra fields; reinterpret as CallToolResult. payload_dict = payload.model_dump(by_alias=True, exclude_none=True) - payload_dict.pop("_meta", None) try: return types.CallToolResult.model_validate(payload_dict) except ValidationError as ex: @@ -2774,6 +3056,7 @@ def __init__( sampling_approval_callback: SamplingApprovalCallback | None = None, sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS, sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, additional_properties: dict[str, Any] | None = None, task_options: MCPTaskOptions | None = None, additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, @@ -2842,6 +3125,8 @@ def __init__( (``min(requested, cap)``); ``None`` disables it. sampling_max_requests: Per-session cap on the number of sampling requests; further requests are rejected. Resets on reconnect. ``None`` disables it. + max_host_payload_size_bytes: Maximum encoded MCP result size retained for Host + transports. ``None`` disables the limit. task_options: Options for tools that advertise ``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`. additional_tool_argument_names: Extra argument names to forward to the MCP server in @@ -2889,6 +3174,7 @@ def __init__( sampling_approval_callback=sampling_approval_callback, sampling_max_tokens=sampling_max_tokens, sampling_max_requests=sampling_max_requests, + max_host_payload_size_bytes=max_host_payload_size_bytes, ) self.command = command self.args = args or [] @@ -2969,6 +3255,7 @@ def __init__( sampling_approval_callback: SamplingApprovalCallback | None = None, sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS, sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, additional_properties: dict[str, Any] | None = None, http_client: AsyncClient | None = None, header_provider: Callable[[dict[str, Any]], dict[str, str]] | None = None, @@ -3038,6 +3325,8 @@ def __init__( (``min(requested, cap)``); ``None`` disables it. sampling_max_requests: Per-session cap on the number of sampling requests; further requests are rejected. Resets on reconnect. ``None`` disables it. + max_host_payload_size_bytes: Maximum encoded MCP result size retained for Host + transports. ``None`` disables the limit. http_client: Optional asyncClient to use. If not provided, the ``streamable_http_client`` API will create and manage a default client. To configure headers, timeouts, or other HTTP client settings, create @@ -3115,6 +3404,7 @@ def __init__( sampling_approval_callback=sampling_approval_callback, sampling_max_tokens=sampling_max_tokens, sampling_max_requests=sampling_max_requests, + max_host_payload_size_bytes=max_host_payload_size_bytes, ) self.url = url self.terminate_on_close = terminate_on_close @@ -3289,6 +3579,7 @@ def __init__( sampling_approval_callback: SamplingApprovalCallback | None = None, sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS, sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, additional_properties: dict[str, Any] | None = None, task_options: MCPTaskOptions | None = None, additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, @@ -3355,6 +3646,8 @@ def __init__( (``min(requested, cap)``); ``None`` disables it. sampling_max_requests: Per-session cap on the number of sampling requests; further requests are rejected. Resets on reconnect. ``None`` disables it. + max_host_payload_size_bytes: Maximum encoded MCP result size retained for Host + transports. ``None`` disables the limit. task_options: Options for tools that advertise ``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`. additional_tool_argument_names: Extra argument names to forward to the MCP server in @@ -3402,6 +3695,7 @@ def __init__( sampling_approval_callback=sampling_approval_callback, sampling_max_tokens=sampling_max_tokens, sampling_max_requests=sampling_max_requests, + max_host_payload_size_bytes=max_host_payload_size_bytes, ) self.url = url self._client_kwargs = kwargs diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 9b407547608..da03ab4e96c 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1433,11 +1433,16 @@ def _function_execution_error_result( message = "Error: Function failed." if config.get("include_detailed_errors", False): message = f"{message} Exception: {exception}" + additional_properties = dict(function_call.additional_properties) + # Tool-specific exceptions can retain caller-only transport metadata without changing the model-facing error. + exception_properties = getattr(exception, "_function_result_additional_properties", None) + if isinstance(exception_properties, Mapping): + additional_properties.update(cast(Mapping[str, Any], exception_properties)) return Content.from_function_result( call_id=function_call.call_id, # type: ignore[arg-type] result=message, exception=str(exception), - additional_properties=function_call.additional_properties, + additional_properties=additional_properties, ) 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", ] diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 669c5d931a6..f1177456afe 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -31,15 +31,22 @@ ) from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning from agent_framework._mcp import ( + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, MCPTool, _build_prefixed_mcp_name, _get_input_model_from_mcp_prompt, + _json_size_exceeds, + _make_mcp_tool_caller, + _mcp_tool_result_host_payload, + _MCPToolResultException, _normalize_additional_tool_argument_names, _normalize_mcp_name, _should_propagate_cancelled_error, + _with_mcp_tool_result_host_payload, logger, ) from agent_framework._middleware import FunctionMiddlewarePipeline +from agent_framework._tools import _function_execution_error_result, normalize_function_invocation_configuration from agent_framework.exceptions import ToolException, ToolExecutionException # Integration test skip condition @@ -60,6 +67,12 @@ def _mcp_result_to_text(result: str | list[Content]) -> str: _HELPER_MCP_TOOL = MCPTool(name="helper") # type: ignore[abstract] +async def _call_generated_mcp_tool(tool: MCPTool, tool_name: str, **kwargs: Any) -> str | list[Content]: + function = FunctionTool(name=tool_name, description="", func=None, input_model={}) + context = FunctionInvocationContext(function=function, arguments=kwargs) + return await _make_mcp_tool_caller(tool, tool_name)(context, **kwargs) + + def _reset_progressive_mcp_warning_state() -> None: _WARNED_FEATURES.discard((ExperimentalWarning, ExperimentalFeature.PROGRESSIVE_TOOLS.value)) @@ -491,6 +504,153 @@ def test_parse_tool_result_from_mcp_structured_content_with_text(): assert parsed == {"data": [1, 2, 3]} +def test_parse_tool_result_from_mcp_preserves_complete_host_payload_once(): + """The complete MCP result survives model parsing and Content persistence.""" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Summary")], + structuredContent={"image_url": "https://example.test/widget.png"}, + isError=False, + _meta={"widget": "image"}, + ) + + model_items = _HELPER_MCP_TOOL._parse_tool_result_from_mcp(mcp_result) + parsed = _with_mcp_tool_result_host_payload(model_items, mcp_result, max_size_bytes=None) + expected_host_payload = { + "_meta": {"widget": "image"}, + "content": [{"type": "text", "text": "Summary"}], + "structuredContent": {"image_url": "https://example.test/widget.png"}, + "isError": False, + } + + assert parsed[0].additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] == expected_host_payload + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in parsed[1].additional_properties + function_result = Content.from_function_result(call_id="call-1", result=parsed) + restored = Content.from_dict(function_result.to_dict()) + + assert restored.result == function_result.result + assert restored.items is not None + assert restored.items[0].additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] == expected_host_payload + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in restored.items[1].additional_properties + + +async def test_custom_mcp_result_parser_preserves_host_payload_and_model_projection() -> None: + """A custom parser controls model content while core retains the complete Host payload once.""" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Server summary")], + structuredContent={"image_url": "https://example.test/widget.png"}, + ) + tool = MCPTool(name="helper", parse_tool_results=lambda _: "Custom model summary") # type: ignore[abstract] + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + direct_result = await tool.call_tool("widget") + parsed = await _call_generated_mcp_tool(tool, "widget") + + assert direct_result == "Custom model summary" + assert isinstance(parsed, list) + assert [item.text for item in parsed] == ["Custom model summary"] + assert parsed[0].additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "image_url": "https://example.test/widget.png" + } + + +async def test_oversized_mcp_host_payload_is_omitted_without_changing_model_result( + caplog: pytest.LogCaptureFixture, +) -> None: + """An oversized Host payload is not retained, while custom model content and metadata survive.""" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Server summary")], + structuredContent={"widget_data": "x" * 1024}, + _meta={"source": "oversized"}, + ) + tool = MCPTool( # type: ignore[abstract] + name="helper", + parse_tool_results=lambda _: "Bounded model summary", + max_host_payload_size_bytes=128, + ) + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + with caplog.at_level(logging.WARNING): + parsed = await _call_generated_mcp_tool(tool, "widget") + + assert isinstance(parsed, list) + assert [item.text for item in parsed] == ["Bounded model summary"] + assert parsed[0].additional_properties["_meta"] == {"source": "oversized"} + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in parsed[0].additional_properties + assert "Omitting MCP Host payload" in caplog.text + + +def test_mcp_host_payload_size_limit_must_be_positive_or_none() -> None: + with pytest.raises(ValueError, match="positive or None"): + MCPTool(name="invalid", max_host_payload_size_bytes=0) # type: ignore[abstract] + + unlimited = MCPTool(name="unlimited", max_host_payload_size_bytes=None) # type: ignore[abstract] + assert unlimited.max_host_payload_size_bytes is None + + +def test_mcp_host_payload_size_preflight_matches_json_and_aborts_before_dump( + monkeypatch: pytest.MonkeyPatch, +) -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text='Escaped "\n\u2603" text')], + structuredContent={"widget_data": "x" * 1024}, + ) + encoded_size = len(json.dumps(mcp_result.model_dump(by_alias=True, exclude_none=True)).encode("utf-8")) + + assert _json_size_exceeds(mcp_result, encoded_size - 1) is True + assert _json_size_exceeds(mcp_result, encoded_size) is False + + uri_result = types.CallToolResult( + content=[ + types.ResourceLink( + type="resource_link", + uri=AnyUrl("file:///abc"), + name="resource", + ) + ] + ) + uri_payload = _mcp_tool_result_host_payload(uri_result, max_size_bytes=None) + assert uri_payload is not None + uri_size = len(json.dumps(uri_payload).encode("utf-8")) + assert _mcp_tool_result_host_payload(uri_result, max_size_bytes=uri_size) == uri_payload + assert _mcp_tool_result_host_payload(uri_result, max_size_bytes=uri_size - 1) is None + + def fail_if_dumped(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("oversized payload must be rejected before model_dump") + + monkeypatch.setattr(types.CallToolResult, "model_dump", fail_if_dumped) + assert _mcp_tool_result_host_payload(mcp_result, max_size_bytes=128) is None + + +async def test_mcp_error_preserves_complete_host_payload_on_function_result(): + """An MCP error keeps its Host payload after generic function error conversion.""" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Widget failed")], + structuredContent={"reason": "invalid input"}, + isError=True, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + with pytest.raises(_MCPToolResultException) as exc_info: + await tool.call_tool("widget") + + function_result = _function_execution_error_result( + Content.from_function_call(call_id="call-1", name="widget"), + "widget", + exc_info.value, + normalize_function_invocation_configuration(None), + ) + host_payload = function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] + + assert function_result.result == "Error: Function failed." + assert host_payload["content"] == [{"type": "text", "text": "Widget failed"}] + assert host_payload["structuredContent"] == {"reason": "invalid input"} + assert host_payload["isError"] is True + + def test_parse_tool_result_from_mcp_structured_content_none(): """Test that None structuredContent does not affect results.""" mcp_result = types.CallToolResult( @@ -6635,11 +6795,21 @@ def _make_create_task_result(task_id: str = "task-1") -> types.CreateTaskResult: ) -def _make_payload(text: str = "done!", is_error: bool = False) -> types.GetTaskPayloadResult: - return types.GetTaskPayloadResult.model_validate({ +def _make_payload( + text: str = "done!", + is_error: bool = False, + structured_content: dict[str, Any] | None = None, + meta: dict[str, Any] | None = None, +) -> types.GetTaskPayloadResult: + payload: dict[str, Any] = { "content": [{"type": "text", "text": text}], "isError": is_error, - }) + } + if structured_content is not None: + payload["structuredContent"] = structured_content + if meta is not None: + payload["_meta"] = meta + return types.GetTaskPayloadResult.model_validate(payload) def _make_task_tool( @@ -6740,22 +6910,58 @@ async def test_call_tool_routes_required_through_task_lifecycle(monkeypatch: pyt monkeypatch.setattr(_mcp_module, "_MCP_TASK_MIN_POLL_INTERVAL", _mcp_module.timedelta(milliseconds=1)) tool = _make_task_tool() + tool.parse_tool_results = lambda _: "custom task summary" 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="working")), ("tasks/get", _make_task_snapshot(status="completed")), - ("tasks/result", _make_payload("hello task")), + ( + "tasks/result", + _make_payload( + "hello task", + structured_content={"widget": "task"}, + meta={"source": "completed-task"}, + ), + ), ) ) - result = await tool.call_tool("slow_op", x=1) + result = await _call_generated_mcp_tool(tool, "slow_op", x=1) - assert _mcp_result_to_text(result) == "hello task" + assert _mcp_result_to_text(result) == "custom task summary" + assert isinstance(result, list) + assert result[0].additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == {"widget": "task"} + assert result[0].additional_properties["_meta"] == {"source": "completed-task"} + assert result[0].additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["_meta"] == {"source": "completed-task"} # Plain session.call_tool must NOT be used for required tools. tool.session.call_tool.assert_not_called() # type: ignore[union-attr] # ty: ignore[unresolved-attribute] +async def test_call_tool_as_task_fallback_preserves_custom_parser_host_payload() -> None: + """A legacy non-task response retains the Host payload after custom parsing.""" + tool = _make_task_tool() + tool.parse_tool_results = lambda _: "custom fallback summary" + fallback_result = types.CallToolResult( + content=[types.TextContent(type="text", text="fallback")], + structuredContent={"widget": "fallback"}, + _meta={"source": "fallback"}, + ) + tool.session.send_request = AsyncMock( # type: ignore[method-assign, union-attr] # ty: ignore[invalid-assignment] + return_value=types.Result.model_validate(fallback_result.model_dump(by_alias=True, exclude_none=True)) + ) + + result = await _call_generated_mcp_tool(tool, "slow_op") + + assert _mcp_result_to_text(result) == "custom fallback summary" + assert isinstance(result, list) + assert result[0].additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": "fallback" + } + assert result[0].additional_properties["_meta"] == {"source": "fallback"} + assert result[0].additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["_meta"] == {"source": "fallback"} + + async def test_call_tool_as_task_default_ttl_propagates() -> None: from datetime import timedelta