diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index f1bf4ea4bc..c652b42d21 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -14,14 +14,16 @@ 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 from opentelemetry import propagate from opentelemetry import trace as otel_trace +from pydantic_core import to_jsonable_python from ._feature_stage import ( ExperimentalFeature, @@ -29,8 +31,15 @@ _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 ._tools import ( + _FUNCTION_RESULT_CARRIER_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] + _FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] + FunctionTool, + _FunctionResultCarrier, # pyright: ignore[reportPrivateUsage] + _FunctionResultPayloadBudget, # pyright: ignore[reportPrivateUsage] +) from ._types import ( ChatOptions, Content, @@ -90,6 +99,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" @@ -128,6 +138,255 @@ class MCPSpecificApproval(TypedDict, total=False): _MCP_INJECTED_HEADER_KEYS_EXTENSION = "agent_framework.mcp_injected_header_keys" MCP_DEFAULT_TIMEOUT = 30 MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5 +_DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES = 1024 * 1024 + + +@dataclass +class _MCPHostPayloadCapture: + """Collect one generated MCP result without affecting its public return shape.""" + + max_size_bytes: int | None + aggregate_budget: _FunctionResultPayloadBudget | None + 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 + return self.meta + + def record(self, mcp_type: Any) -> None: + if self.recorded: + return + self.recorded = True + self.prepare_meta(mcp_type) + + 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 Host payload because the request retention budget is exhausted.") + return + if remaining is not None: + effective_limit = min(effective_limit, remaining) if effective_limit is not None else remaining + + host_payload = _mcp_tool_result_host_payload(mcp_type, max_size_bytes=effective_limit) + if host_payload is None: + return + encoded_size = len(json.dumps(host_payload).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 Host payload because the request retention budget is exhausted.") + return + self.host_payload = host_payload + + def to_carrier(self) -> _FunctionResultCarrier: + additional_properties: dict[str, Any] = {} + if self.meta is not None: + additional_properties["_meta"] = self.meta + if self.host_payload is not None: + additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] = self.host_payload + return _FunctionResultCarrier( + additional_properties=additional_properties, + item_additional_properties={"_meta": self.meta} if self.meta is not None else {}, + exclusive_outer_keys=frozenset({_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY}), + exclusive_item_keys=frozenset({"_meta"}), + result_already_parsed=True, + ) + + def prepare_model_result(self, parsed: str | list[Content]) -> list[Content]: + """Apply bounded metadata before security middleware sees generated output.""" + items = [Content.from_text(parsed)] if isinstance(parsed, str) else list(parsed) + for index, item in enumerate(items): + 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["_meta"] = self.meta + else: + updated_item.additional_properties.pop("_meta", None) + items[index] = updated_item + return items + + +_mcp_host_payload_capture: contextvars.ContextVar[_MCPHostPayloadCapture | None] = contextvars.ContextVar( + "_mcp_host_payload_capture", + default=None, +) + + +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: + budget.consume(6) + elif codepoint > 0xFFFF: + budget.consume(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 callable(getattr(value, "unicode_string", None)): + _consume_json_string(str(value), budget) + return + if isinstance(value, Mapping): + budget.consume(2) + for index, (key, item) in enumerate(cast(Mapping[Any, Any], value).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) + for index, item in enumerate(cast(Sequence[Any], value)): + 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): + for key, item in cast(Mapping[Any, Any], model_extra_raw).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 bounded, JSON-safe copy of a complete MCP result.""" + 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 _mcp_tool_result_meta( + mcp_type: Any, + *, + max_size_bytes: int | None, +) -> dict[str, Any] | None: + """Return a separately bounded, JSON-safe copy of MCP result metadata.""" + raw_meta = getattr(mcp_type, "meta", None) + if not isinstance(raw_meta, Mapping): + return None + if max_size_bytes is not None and _json_size_exceeds(raw_meta, max_size_bytes): + logger.warning( + "Omitting MCP result _meta because its encoded size exceeds max_host_payload_size_bytes=%d.", + max_size_bytes, + ) + return None + meta = cast(dict[str, Any], to_jsonable_python(dict(cast(Mapping[str, Any], raw_meta)), fallback=str)) + if max_size_bytes is not None and len(json.dumps(meta).encode("utf-8")) > max_size_bytes: + logger.warning( + "Omitting MCP result _meta because its encoded size exceeds max_host_payload_size_bytes=%d.", + max_size_bytes, + ) + return None + return meta + + +def _capture_mcp_tool_result(mcp_type: Any) -> None: + capture = _mcp_host_payload_capture.get() + if capture is not None: + capture.record(mcp_type) class _MCPHeaderScopedClient: @@ -297,7 +556,18 @@ 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) + raw_budget = ctx.metadata.get(_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY) + capture = _MCPHostPayloadCapture( + max_size_bytes=mcp_tool.max_host_payload_size_bytes, + aggregate_budget=raw_budget if isinstance(raw_budget, _FunctionResultPayloadBudget) else None, + ) + token = _mcp_host_payload_capture.set(capture) + try: + parsed = await mcp_tool.call_tool(remote_tool_name, **call_kwargs) + return capture.prepare_model_result(parsed) + finally: + _mcp_host_payload_capture.reset(token) + ctx.metadata[_FUNCTION_RESULT_CARRIER_CONTEXT_KEY] = capture.to_carrier() return _call_tool_with_runtime_kwargs @@ -541,6 +811,8 @@ def __init__( additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, use_progressive_disclosure: bool = False, always_load: Collection[str] | None = None, + *, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, ) -> None: """Initialize the MCP Tool base. @@ -607,6 +879,9 @@ def __init__( always_load: MCP tool names to keep visible from the start when progressive disclosure is enabled. Names use the same safe matching rules as ``allowed_tools``; unmatched entries are ignored. + 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. """ if use_progressive_disclosure and not load_tools: raise ValueError("use_progressive_disclosure=True requires load_tools=True.") @@ -617,6 +892,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 @@ -627,6 +904,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 @@ -750,13 +1028,19 @@ 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. + + Generated MCP functions also preserve the complete MCP result under a + private core-owned ``additional_properties`` marker for Host transports. + This does not change which content is 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. + capture = _mcp_host_payload_capture.get() + if capture is not None: + meta = capture.prepare_meta(mcp_type) + else: + raw_meta = mcp_type.meta + meta = dict(raw_meta) if isinstance(raw_meta, Mapping) else None additional_kwargs: dict[str, Any] = {"additional_properties": {"_meta": meta}} if meta else {} result: list[Content] = [] @@ -801,7 +1085,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)) @@ -2224,6 +2508,13 @@ 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) + + async def _call_tool( + self, + tool_name: str, + kwargs: dict[str, Any], + ) -> 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." @@ -2245,7 +2536,13 @@ 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, + ) async def _call_tool_with_retries( self, @@ -2264,6 +2561,7 @@ async def _call_tool_with_retries( result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=meta) # type: ignore 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) @@ -2273,7 +2571,9 @@ 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)) - return parser(result) + parsed = parser(result) + _capture_mcp_tool_result(result) + return parsed except ToolExecutionException: raise except (ClosedResourceError, McpError) as call_ex: @@ -2383,6 +2683,13 @@ 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) + + async def _call_tool_as_task( + self, + tool_name: str, + kwargs: dict[str, Any], + ) -> str | list[Content]: from anyio import ClosedResourceError from mcp.shared.exceptions import McpError @@ -2417,13 +2724,16 @@ async def call_tool_as_task(self, tool_name: str, **kwargs: Any) -> str | list[C if fallback_result is not None: 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)) - return parser(fallback_result) + parsed = parser(fallback_result) + _capture_mcp_tool_result(fallback_result) + return parsed if task_id is None: raise ToolExecutionException(f"MCP server did not return a task_id or fallback result for '{tool_name}'.") @@ -2435,7 +2745,12 @@ 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, + ) try: if max_wait_s is not None: @@ -2515,7 +2830,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): @@ -2612,13 +2926,16 @@ async def _handle_terminal_task( payload = await self._fetch_task_result(task_id) 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)) - return parser(payload) + parsed = parser(payload) + _capture_mcp_tool_result(payload) + return parsed # Non-completed terminal statuses surface as ToolExecutionException so the # function-calling loop sees a normal failure for tool_name. @@ -2650,7 +2967,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: @@ -2931,6 +3247,7 @@ def __init__( additional_properties: dict[str, Any] | None = None, task_options: MCPTaskOptions | None = None, additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, **kwargs: Any, ) -> None: """Initialize the MCP stdio tool. @@ -3021,6 +3338,8 @@ def __init__( Treat those keys as visible to this server, and source credentials outside ``function_invocation_kwargs`` - for example through ``env`` - for servers whose process you do not control. + max_host_payload_size_bytes: Maximum encoded MCP result size retained for Host + transports. ``None`` disables the limit. kwargs: Any extra arguments to pass to the stdio client. """ super().__init__( @@ -3044,6 +3363,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 [] @@ -3129,6 +3449,7 @@ def __init__( header_provider: Callable[[dict[str, Any]], dict[str, str]] | None = None, task_options: MCPTaskOptions | None = None, additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, **kwargs: Any, ) -> None: """Initialize the MCP streamable HTTP tool. @@ -3251,6 +3572,8 @@ def __init__( ``function_invocation_kwargs``: read a ``ContextVar`` inside the provider (which still allows a different value per request), or configure a custom ``http_client``. + max_host_payload_size_bytes: Maximum encoded MCP result size retained for Host + transports. ``None`` disables the limit. kwargs: Additional keyword arguments (accepted for backward compatibility but not used). """ super().__init__( @@ -3274,6 +3597,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 @@ -3500,6 +3824,7 @@ def __init__( additional_properties: dict[str, Any] | None = None, task_options: MCPTaskOptions | None = None, additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None, + max_host_payload_size_bytes: int | None = _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES, **kwargs: Any, ) -> None: """Initialize the MCP WebSocket tool. @@ -3588,6 +3913,8 @@ def __init__( The same dict is shared with every MCP server attached to the run. This transport has no header hook, so source credentials outside ``function_invocation_kwargs`` for servers you do not control. + max_host_payload_size_bytes: Maximum encoded MCP result size retained for Host + transports. ``None`` disables the limit. kwargs: Any extra arguments to pass to the WebSocket client. """ super().__init__( @@ -3611,6 +3938,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 ac3fc903d9..3b298b0863 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -106,6 +106,9 @@ def _generate_function_call_occurrence_id() -> str: _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY: Final[str] = "already_approved_approval_request_groups" _PENDING_APPROVAL_REQUESTS_KEY: Final[str] = "pending_approval_requests" _FUNCTION_INVOCATION_BUDGET_STATE_KEY: Final[str] = "_function_invocation_budget_state" +_FUNCTION_RESULT_CARRIER_CONTEXT_KEY: Final[str] = "_function_result_carrier" +_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY: Final[str] = "_function_result_payload_budget" +_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY: Final[str] = "_function_result_payload_budget" _FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT: Final[str] = ( "Function invocation limit reached before a final answer could be produced." ) @@ -115,6 +118,40 @@ def _generate_function_call_occurrence_id() -> str: ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) +@dataclass +class _FunctionResultCarrier: + """Host-only properties retained independently of model-facing tool output.""" + + additional_properties: dict[str, Any] + item_additional_properties: dict[str, Any] + exclusive_outer_keys: frozenset[str] = frozenset() + exclusive_item_keys: frozenset[str] = frozenset() + 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 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) + return max(self.limit_bytes - self.retained_bytes, 0) + + def reserve(self, size_bytes: int, per_result_limit: int | None) -> bool: + if per_result_limit is None: + return True + remaining = self.remaining(per_result_limit) + if remaining is None or size_bytes > remaining: + return False + self.retained_bytes += size_bytes + return True + + class _SkipParsingSentinel: """Sentinel signaling that :meth:`FunctionTool.invoke` should return the raw value. @@ -722,11 +759,21 @@ async def invoke( logger.info(f"Function {self.name} succeeded.") logger.debug(f"Function result: {type(result).__name__}") return result - try: - parsed = parser(result) - except Exception: - logger.warning(f"Function {self.name}: result parser failed, falling back to str().") - parsed = [Content.from_text(str(result))] + carrier = ( + effective_context.metadata.get(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY) if effective_context else None + ) + if ( + isinstance(carrier, _FunctionResultCarrier) + and carrier.result_already_parsed + and configured_parser is None + ): + parsed = result + else: + try: + parsed = parser(result) + except Exception: + logger.warning(f"Function {self.name}: result parser failed, falling back to str().") + parsed = [Content.from_text(str(result))] if isinstance(parsed, str): parsed = [Content.from_text(parsed)] logger.info(f"Function {self.name} succeeded.") @@ -787,11 +834,21 @@ async def invoke( if emit_tool_call_attrs: span.set_attribute(OtelAttr.TOOL_RESULT, result_str) return result - try: - parsed = parser(result) - except Exception: - logger.warning(f"Function {self.name}: result parser failed, falling back to str().") - parsed = [Content.from_text(str(result))] + carrier = ( + effective_context.metadata.get(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY) if effective_context else None + ) + if ( + isinstance(carrier, _FunctionResultCarrier) + and carrier.result_already_parsed + and configured_parser is None + ): + parsed = result + else: + try: + parsed = parser(result) + except Exception: + logger.warning(f"Function {self.name}: result parser failed, falling back to str().") + parsed = [Content.from_text(str(result))] if isinstance(parsed, str): parsed = [Content.from_text(parsed)] logger.info(f"Function {self.name} succeeded.") @@ -1440,9 +1497,8 @@ def _function_execution_error_result( tool_name: str, exception: Exception, config: FunctionInvocationConfiguration, + context: FunctionInvocationContext | None = None, ) -> Content: - from ._types import Content - logger.warning( "Function '%s' raised an exception; returning an error result to the model. " "Set include_detailed_errors=True for the full detail. Exception: %r", @@ -1452,12 +1508,57 @@ def _function_execution_error_result( message = "Error: Function failed." if config.get("include_detailed_errors", False): message = f"{message} Exception: {exception}" - return Content.from_function_result( + return _finalize_function_result( call_id=function_call.call_id, # type: ignore[arg-type] result=message, exception=str(exception), - additional_properties=function_call.additional_properties, + base_additional_properties=function_call.additional_properties, + context=context, + ) + + +def _finalize_function_result( + *, + call_id: str, + result: Any, + base_additional_properties: Mapping[str, Any] | None = None, + exception: str | None = None, + context: FunctionInvocationContext | None = None, +) -> Content: + """Build the stable function-result wrapper and apply private Host metadata.""" + from ._types import Content + + carrier: _FunctionResultCarrier | None = None + if context is not None: + raw_carrier = context.metadata.pop(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY, None) + if isinstance(raw_carrier, _FunctionResultCarrier): + carrier = raw_carrier + + additional_properties = dict(base_additional_properties or {}) + if carrier is not None: + for key in carrier.exclusive_outer_keys: + additional_properties.pop(key, None) + additional_properties.update(carrier.additional_properties) + + function_result = Content.from_function_result( + call_id=call_id, + result=result, + exception=exception, + additional_properties=additional_properties, ) + if carrier is None or function_result.items is None: + return function_result + + updated_items = list(function_result.items) + for index, item in enumerate(updated_items): + updated_item = copy.copy(item) + updated_item.additional_properties = dict(updated_item.additional_properties) + for key in carrier.exclusive_outer_keys | carrier.exclusive_item_keys: + updated_item.additional_properties.pop(key, None) + updated_item.additional_properties.update(carrier.item_additional_properties) + updated_items[index] = updated_item + function_result.items = updated_items + return function_result async def _auto_invoke_function( @@ -1469,6 +1570,7 @@ async def _auto_invoke_function( invocation_session: AgentSession | None = None, middleware_pipeline: FunctionMiddlewarePipeline | None = None, live_tools: list[ToolTypes] | None = None, + host_payload_budget: _FunctionResultPayloadBudget | None = None, ) -> Content: """Invoke a function call requested by the agent, applying middleware that is defined. @@ -1483,6 +1585,7 @@ async def _auto_invoke_function( middleware_pipeline: Optional middleware pipeline to apply during execution. live_tools: The live, mutable tools list for the current agent run, exposed on the FunctionInvocationContext so tools can add/remove tools at runtime. + host_payload_budget: Shared request budget for retained Host-only function result payloads. Returns: The function result content. @@ -1572,8 +1675,8 @@ async def _auto_invoke_function( if middleware_pipeline is None or not middleware_pipeline.has_middlewares: # No middleware - execute directly + direct_context = None try: - direct_context = None if getattr(tool, "_context_parameter_name", None): direct_context = FunctionInvocationContext( function=tool, @@ -1582,22 +1685,25 @@ async def _auto_invoke_function( kwargs=runtime_kwargs.copy(), tools=live_tools, ) + if host_payload_budget is not None: + direct_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget function_result = await tool.invoke( arguments=args, context=direct_context, tool_call_id=function_call_content.call_id, ) - return Content.from_function_result( + return _finalize_function_result( call_id=function_call_content.call_id, # type: ignore[arg-type] result=function_result, - additional_properties=function_call_content.additional_properties, + base_additional_properties=function_call_content.additional_properties, + context=direct_context, ) except (MiddlewareFailure, UserInputRequiredException): # Explicit control-flow signals escape the loop; only ordinary exceptions # are absorbed into tool-error results below. raise except Exception as exc: - return _function_execution_error_result(function_call_content, tool.name, exc, config) + return _function_execution_error_result(function_call_content, tool.name, exc, config, direct_context) # Execute through middleware pipeline if available middleware_context = FunctionInvocationContext( function=tool, @@ -1606,6 +1712,8 @@ async def _auto_invoke_function( kwargs=runtime_kwargs.copy(), tools=live_tools, ) + if host_payload_budget is not None: + middleware_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget call_id = function_call_content.call_id if call_id is None: @@ -1641,7 +1749,12 @@ async def final_function_handler(context_obj: Any) -> Any: if isinstance(function_result, Content) and function_result.type == "function_approval_request": return function_result - return Content.from_function_result(call_id=call_id, result=function_result) + return _finalize_function_result( + call_id=call_id, + result=function_result, + base_additional_properties=function_call_content.additional_properties, + context=middleware_context, + ) except MiddlewareTermination as term_exc: # Re-raise to signal loop termination, but first capture any result set by middleware if middleware_context.result is not None: @@ -1654,10 +1767,11 @@ async def final_function_handler(context_obj: Any) -> Any: term_exc.result = middleware_context.result else: # Store result in exception for caller to extract - term_exc.result = Content.from_function_result( + term_exc.result = _finalize_function_result( call_id=call_id, result=middleware_context.result, - additional_properties=function_call_content.additional_properties, + base_additional_properties=function_call_content.additional_properties, + context=middleware_context, ) raise except (MiddlewareFailure, UserInputRequiredException): @@ -1666,7 +1780,7 @@ async def final_function_handler(context_obj: Any) -> Any: # relying on the tool-error conversion below, and it propagates to the caller. raise except Exception as exc: - return _function_execution_error_result(function_call_content, tool.name, exc, config) + return _function_execution_error_result(function_call_content, tool.name, exc, config, middleware_context) def _get_tool_map( @@ -1698,6 +1812,7 @@ async def _execute_single_function_call( invocation_session: AgentSession | None, middleware_pipeline: FunctionMiddlewarePipeline | None, live_tools: list[ToolTypes] | None, + host_payload_budget: _FunctionResultPayloadBudget | None, ) -> tuple[list[Content], bool]: from ._middleware import MiddlewareTermination from ._sessions import _suspend_run_persistence_gate # pyright: ignore[reportPrivateUsage] @@ -1719,6 +1834,7 @@ async def _execute_single_function_call( middleware_pipeline=middleware_pipeline, config=config, live_tools=live_tools, + host_payload_budget=host_payload_budget, ) return [result], False except MiddlewareTermination as exc: @@ -1757,6 +1873,7 @@ async def _try_execute_function_call_groups( config: FunctionInvocationConfiguration, invocation_session: AgentSession | None = None, middleware_pipeline: FunctionMiddlewarePipeline | None = None, + host_payload_budget: _FunctionResultPayloadBudget | None = None, ) -> tuple[list[list[Content]], bool]: """Execute multiple function calls concurrently while preserving per-call result groups. @@ -1767,6 +1884,7 @@ async def _try_execute_function_call_groups( config: Configuration for function invocation. invocation_session: The agent session for this invocation, if any. middleware_pipeline: Optional middleware pipeline to apply during execution. + host_payload_budget: Shared request budget for retained Host-only function result payloads. Returns: A tuple of: @@ -1885,6 +2003,7 @@ async def _try_execute_function_call_groups( invocation_session=invocation_session, middleware_pipeline=middleware_pipeline, live_tools=live_tools, + host_payload_budget=host_payload_budget, ), ) for function_call in function_calls @@ -1955,6 +2074,7 @@ async def _execute_function_calls( config: FunctionInvocationConfiguration, invocation_session: AgentSession | None = None, middleware_pipeline: FunctionMiddlewarePipeline | None = None, + host_payload_budget: _FunctionResultPayloadBudget | None = None, ) -> _FunctionExecutionBatch: tools = _extract_tools(options) if not tools: @@ -1966,6 +2086,7 @@ async def _execute_function_calls( invocation_session=invocation_session, middleware_pipeline=middleware_pipeline, config=config, + host_payload_budget=host_payload_budget, ) return _FunctionExecutionBatch( result_groups=result_groups, @@ -3858,6 +3979,13 @@ def get_response( # setdefault preserves the original timestamp across approval re-entries so that # 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() + ) + budget_state[_FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY] = host_payload_budget max_errors = self.function_invocation_configuration.get( "max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST ) @@ -3879,6 +4007,7 @@ def get_response( config=self.function_invocation_configuration, invocation_session=invocation_session, middleware_pipeline=function_middleware_pipeline, + host_payload_budget=host_payload_budget, ) # Give the loop private mutable options and one shared run-local tool list for progressive tool changes. diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index f93727edb2..0784fc53a2 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -2,6 +2,7 @@ # pyright: ignore[reportPrivateUsage] import asyncio import contextlib +import inspect import json import logging import os @@ -20,6 +21,8 @@ from pydantic import AnyUrl, BaseModel from agent_framework import ( + ChatResponse, + ChatResponseUpdate, Content, FunctionInvocationContext, FunctionMiddleware, @@ -32,16 +35,26 @@ from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning from agent_framework._mcp import ( _MCP_HEADER_OWNER_EXTENSION, + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, MCPTool, _build_prefixed_mcp_name, _describe_error, _get_input_model_from_mcp_prompt, + _json_size_exceeds, + _make_mcp_tool_caller, + _mcp_tool_result_host_payload, + _mcp_tool_result_meta, _normalize_additional_tool_argument_names, _normalize_mcp_name, _should_propagate_cancelled_error, logger, ) from agent_framework._middleware import FunctionMiddlewarePipeline +from agent_framework._tools import ( + _auto_invoke_function, + _FunctionResultPayloadBudget, + normalize_function_invocation_configuration, +) from agent_framework.exceptions import ToolException, ToolExecutionException # Integration test skip condition @@ -62,6 +75,34 @@ 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, + *, + result_parser: Any = None, + middleware_pipeline: FunctionMiddlewarePipeline | None = None, + host_payload_budget: _FunctionResultPayloadBudget | None = None, + **kwargs: Any, +) -> Content: + function_kwargs: dict[str, Any] = {} + if result_parser is not None: + function_kwargs["result_parser"] = result_parser + function = FunctionTool( + name=tool_name, + description="", + func=_make_mcp_tool_caller(tool, tool_name), + input_model={"type": "object", "properties": {name: {} for name in kwargs}}, + **function_kwargs, + ) + return await _auto_invoke_function( + Content.from_function_call(call_id=f"call-{tool_name}", name=tool_name, arguments=kwargs), + config=normalize_function_invocation_configuration(None), + tool_map={tool_name: function}, + middleware_pipeline=middleware_pipeline, + host_payload_budget=host_payload_budget, + ) + + def _reset_progressive_mcp_warning_state() -> None: _WARNED_FEATURES.discard((ExperimentalWarning, ExperimentalFeature.PROGRESSIVE_TOOLS.value)) @@ -503,6 +544,512 @@ def test_parse_tool_result_from_mcp_structured_content_with_text(): assert parsed == {"data": [1, 2, 3]} +async def test_generated_mcp_tool_preserves_complete_host_payload_once() -> None: + """The generated FunctionTool path retains one complete, persistent Host payload.""" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Summary")], + structuredContent={"image_url": "https://example.test/widget.png"}, + isError=False, + _meta={"widget": "image"}, + ) + tool = MCPTool(name="helper") # 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.items is not None + expected_host_payload = { + "_meta": {"widget": "image"}, + "content": [{"type": "text", "text": "Summary"}], + "structuredContent": {"image_url": "https://example.test/widget.png"}, + "isError": False, + } + + assert [item.additional_properties["_meta"] for item in function_result.items] == [{"widget": "image"}] * 2 + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] == expected_host_payload + assert all(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in item.additional_properties for item in function_result.items) + restored = Content.from_dict(function_result.to_dict()) + + assert restored.result == function_result.result + assert restored.items is not None + assert restored.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] == expected_host_payload + assert all(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in item.additional_properties for item in restored.items) + + +async def test_generated_mcp_host_payload_replaces_duplicate_private_markers() -> None: + """Only core's outer complete payload marker survives.""" + stale = {_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: {"stale": True}} + mcp_result = types.CallToolResult(content=[types.TextContent(type="text", text="current")]) + tool = MCPTool( # type: ignore[abstract] + name="helper", + parse_tool_results=lambda _: [ + Content.from_text("one", additional_properties=stale), + Content.from_text("two", additional_properties=stale), + ], + ) + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + function_result = await _call_generated_mcp_tool(tool, "widget") + + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["content"] == [ + {"type": "text", "text": "current"} + ] + assert function_result.items is not None + assert all(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in item.additional_properties for item in function_result.items) + + +async def test_custom_mcp_result_parser_preserves_direct_shape_and_generated_host_payload() -> None: + """A custom parser controls model content while generated calls retain the Host payload.""" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Server summary")], + structuredContent={"image_url": "https://example.test/widget.png"}, + _meta={"source": "server"}, + ) + 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") + function_result = await _call_generated_mcp_tool(tool, "widget") + + 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 function_result.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 omitted while bounded 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): + function_result = await _call_generated_mcp_tool(tool, "widget") + + 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 _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in function_result.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_boundary_uri_serialization_and_early_abort( + 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 + assert uri_payload["content"][0]["uri"] == "file:///abc" + 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_generated_mcp_error_preserves_complete_host_payload_on_function_result() -> None: + """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, + _meta={"source": "server"}, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + function_result = await _call_generated_mcp_tool(tool, "widget") + host_payload = function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] + + assert function_result.result == "Error: Function failed." + assert function_result.additional_properties["_meta"] == {"source": "server"} + assert host_payload["content"] == [{"type": "text", "text": "Widget failed"}] + assert host_payload["structuredContent"] == {"reason": "invalid input"} + assert host_payload["isError"] is True + + +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")]) + error = types.CallToolResult(content=[types.TextContent(type="text", text="failed")], isError=True) + tool = MCPTool( # type: ignore[abstract] + name="helper", + parse_tool_results=lambda result: cast(types.TextContent, result.content[0]).text, + ) + tool.session = Mock() + tool.session.call_tool = AsyncMock(side_effect=[success, error]) + + def fail_if_captured(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("direct calls 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) + + assert await tool.call_tool("widget") == "ok" + with pytest.raises(ToolExecutionException, match="failed"): + 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")], + structuredContent={"widget": "complete"}, + _meta={"source": "server"}, + ) + tool = MCPTool(name="helper", parse_tool_results=lambda _: "MCP parser projection") # type: ignore[abstract] + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + result_parser=lambda _: "FunctionTool parser projection", + ) + + assert function_result.result == "FunctionTool parser projection" + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": "complete" + } + assert function_result.items is not None + assert function_result.items[0].additional_properties["_meta"] == {"source": "server"} + + +@pytest.mark.parametrize("parser_layer", ["mcp", "function"]) +async def test_empty_custom_parser_projection_remains_empty(parser_layer: str) -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="server projection")], + structuredContent={"widget": "complete"}, + ) + tool = MCPTool( # type: ignore[abstract] + name="helper", + parse_tool_results=(lambda _: []) if parser_layer == "mcp" else None, + ) + tool.session = Mock() + tool.session.call_tool = AsyncMock(return_value=mcp_result) + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + result_parser=(lambda _: []) if parser_layer == "function" else None, + ) + + assert function_result.result == "" + assert function_result.items == [] + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": "complete" + } + + +async def test_oversized_mcp_error_preserves_independently_bounded_meta() -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="failed")], + structuredContent={"large": "x" * 1024}, + isError=True, + _meta={"source": "small"}, + ) + tool = MCPTool(name="helper", max_host_payload_size_bytes=128) # 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.exception is not None + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in function_result.additional_properties + assert function_result.additional_properties["_meta"] == {"source": "small"} + assert function_result.items is not None + assert function_result.items[0].additional_properties["_meta"] == {"source": "small"} + + +def test_mcp_result_meta_has_independent_boundary_and_early_abort(monkeypatch: pytest.MonkeyPatch) -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="ok")], + _meta={"source": "server", "uri": AnyUrl("https://example.test/resource")}, + ) + expected = {"source": "server", "uri": "https://example.test/resource"} + encoded_size = len(json.dumps(expected).encode("utf-8")) + + assert _mcp_tool_result_meta(mcp_result, max_size_bytes=encoded_size) == expected + assert _mcp_tool_result_meta(mcp_result, max_size_bytes=encoded_size - 1) is None + + oversized = types.CallToolResult( + content=[types.TextContent(type="text", text="ok")], + _meta={"large": "x" * 1024}, + ) + + def fail_if_copied(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("oversized _meta must be rejected before JSON-safe copying") + + monkeypatch.setattr("agent_framework._mcp.to_jsonable_python", fail_if_copied) + assert _mcp_tool_result_meta(oversized, max_size_bytes=128) is None + + +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")], + _meta={"large": "x" * 1024}, + ) + tool = MCPTool(name="helper", max_host_payload_size_bytes=128) # 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 == "ok" + assert "_meta" not in function_result.additional_properties + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in function_result.additional_properties + assert function_result.items is not None + assert "_meta" not in function_result.items[0].additional_properties + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +async def test_mcp_host_payload_survives_real_function_loop( + chat_client_base: Any, + streaming: bool, +) -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="model projection")], + structuredContent={"widget": "complete"}, + _meta={"source": "server"}, + ) + 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": {}}, + ) + function_call = Content.from_function_call(call_id="call-widget", name="widget", arguments={}) + captured_model_inputs: list[list[Message]] = [] + + if streaming: + original_streaming = chat_client_base._get_streaming_response + + def capture_streaming(*, messages: list[Message], **kwargs: Any) -> Any: + captured_model_inputs.append([Message.from_dict(message.to_dict()) for message in messages]) + return original_streaming(messages=messages, **kwargs) + + chat_client_base._get_streaming_response = capture_streaming + chat_client_base.streaming_responses = [ + [ChatResponseUpdate(role="assistant", contents=[function_call], finish_reason="tool_calls")], + [ + ChatResponseUpdate( + role="assistant", + contents=[Content.from_text("done")], + finish_reason="stop", + ) + ], + ] + stream = chat_client_base.get_response( + [Message(role="user", contents=["run"])], + stream=True, + options={"tools": [function]}, + ) + updates = [update async for update in stream] + response = await stream.get_final_response() + streamed_results = [ + content for update in updates for content in update.contents if content.type == "function_result" + ] + assert len(streamed_results) == 1 + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY in streamed_results[0].additional_properties + else: + original_non_streaming = chat_client_base._get_non_streaming_response + + async def capture_non_streaming(*, messages: list[Message], **kwargs: Any) -> ChatResponse: + captured_model_inputs.append([Message.from_dict(message.to_dict()) for message in messages]) + return await original_non_streaming(messages=messages, **kwargs) + + chat_client_base._get_non_streaming_response = capture_non_streaming + chat_client_base.run_responses = [ + ChatResponse(messages=Message(role="assistant", contents=[function_call])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + response = await chat_client_base.get_response( + [Message(role="user", contents=["run"])], + options={"tools": [function]}, + ) + + assert len(captured_model_inputs) == 2 + second_call_results = [ + content + for message in captured_model_inputs[1] + for content in message.contents + if content.type == "function_result" + ] + final_results = [ + content for message in response.messages for content in message.contents if content.type == "function_result" + ] + assert len(second_call_results) == 1 + assert len(final_results) == 1 + for function_result in [second_call_results[0], final_results[0]]: + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": "complete" + } + assert function_result.items is not None + assert all( + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in item.additional_properties for item in function_result.items + ) + restored = Content.from_dict(function_result.to_dict()) + assert restored.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": "complete" + } + + +@pytest.mark.parametrize( + ("size_limit", "expected_markers"), + [(512, 1), (None, 2)], + ids=["bounded", "unlimited"], +) +async def test_mcp_host_payload_has_aggregate_request_budget( + chat_client_base: Any, + size_limit: int | None, + expected_markers: int, +) -> None: + tool = MCPTool(name="helper", max_host_payload_size_bytes=size_limit) # type: ignore[abstract] + tool.session = Mock() + + 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}, + ) + + tool.session.call_tool = AsyncMock(side_effect=call_tool) + functions = [ + FunctionTool( + name=name, + description="", + func=_make_mcp_tool_caller(tool, name), + input_model={"type": "object", "properties": {}}, + ) + for name in ("one", "two") + ] + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call(call_id=f"call-{function.name}", name=function.name, arguments={}) + for function in functions + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + response = await chat_client_base.get_response( + [Message(role="user", contents=["run"])], + options={"tools": functions}, + ) + + function_results = [ + content for message in response.messages for content in message.contents if content.type == "function_result" + ] + retained_payloads = [ + result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY] + for result in function_results + if _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY in result.additional_properties + ] + 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 + + +async def test_secure_mcp_auto_hide_preserves_outer_host_payload() -> None: + from agent_framework.security import ( + IntegrityLabel, + LabelTrackingFunctionMiddleware, + _wrap_mcp_function_for_ifc, + ) + + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="untrusted payload")], + structuredContent={"widget": "complete"}, + _meta={"ifc": {"integrity": "untrusted", "confidentiality": "public"}}, + ) + 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( + 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) + middleware = LabelTrackingFunctionMiddleware(auto_hide_untrusted=True) + + 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(middleware), + host_payload_budget=_FunctionResultPayloadBudget(), + ) + + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "widget": "complete" + } + 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 hidden_item.text != "untrusted payload" + + def test_parse_tool_result_from_mcp_structured_content_none(): """Test that None structuredContent does not affect results.""" mcp_result = types.CallToolResult( @@ -1757,6 +2304,16 @@ def test_mcp_transport_subclasses_accept_progressive_disclosure_options() -> Non assert websocket.always_load == ["search"] +def test_mcp_transport_subclasses_forward_host_payload_limit() -> None: + tools = [ + MCPStdioTool(name="stdio", command="python", max_host_payload_size_bytes=101), + MCPStreamableHTTPTool(name="http", url="https://example.com/mcp", max_host_payload_size_bytes=102), + MCPWebsocketTool(name="ws", url="wss://example.com/mcp", max_host_payload_size_bytes=103), + ] + + assert [tool.max_host_payload_size_bytes for tool in tools] == [101, 102, 103] + + def test_mcp_progressive_disclosure_requires_loading_tools() -> None: with pytest.raises(ValueError, match="requires load_tools=True"): MCPTool( # type: ignore[abstract] @@ -1775,9 +2332,12 @@ def test_mcp_progressive_disclosure_warns_on_construction() -> None: def test_mcp_tool_base_constructor_preserves_positional_tool_name_prefix() -> None: tool = MCPTool("test_server", "description", None, None, "prefix") # type: ignore[abstract] + parameters = inspect.signature(MCPTool.__init__).parameters assert tool.tool_name_prefix == "prefix" assert tool.use_progressive_disclosure is False + assert parameters["always_load"].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD + assert parameters["max_host_payload_size_bytes"].kind is inspect.Parameter.KEYWORD_ONLY def _progressive_tool_list_page(*, tools: list[types.Tool] | None = None) -> types.ListToolsResult: @@ -6770,19 +7330,24 @@ def provider(kwargs): await server.load_tools() func = server.functions[0] - # Build a FunctionInvocationContext with runtime kwargs, as the agent framework would - context = FunctionInvocationContext( - function=func, - arguments={"name": "Alice"}, - kwargs={"some_token": "my-secret"}, - ) - with patch.object(MCPStreamableHTTPTool, "call_tool", spy_call_tool): - result = await func.invoke(arguments={"name": "Alice"}, context=context) + result = await _auto_invoke_function( + Content.from_function_call( + call_id="call-greet", + name=func.name, + arguments={"name": "Alice"}, + ), + custom_args={"some_token": "my-secret"}, + config=normalize_function_invocation_configuration(None), + tool_map={func.name: func}, + host_payload_budget=_FunctionResultPayloadBudget(), + ) # Verify the invoke produced a result - assert isinstance(result, list) - assert result[0].text == "Hello!" + assert result.items is not None + assert result.items[0].text == "Hello!" + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY in result.additional_properties + assert all(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in item.additional_properties for item in result.items) # Verify header_provider was called with the runtime kwargs assert len(provider_received) == 1 @@ -7017,11 +7582,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( @@ -7122,22 +7697,87 @@ 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) + function_result = await _call_generated_mcp_tool(tool, "slow_op", x=1) - assert _mcp_result_to_text(result) == "hello task" + assert function_result.result == "custom task summary" + assert function_result.items is not None + 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 function_result.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_routes_required_through_public_task_override() -> None: + class OverriddenTaskTool(MCPTool): + def __init__(self) -> None: + super().__init__(name="override") + self.override_called = False + + async def call_tool_as_task(self, tool_name: str, **kwargs: Any) -> str | list[Content]: + self.override_called = True + return await super().call_tool_as_task(tool_name, **kwargs) + + tool = OverriddenTaskTool() # type: ignore[abstract] + tool.session = AsyncMock(spec=ClientSession) + tool._tool_task_support_by_name["slow_op"] = "required" + fallback_result = types.CallToolResult(content=[types.TextContent(type="text", text="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)) + ) + + function_result = await _call_generated_mcp_tool(tool, "slow_op") + + assert function_result.result == "fallback" + assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY in function_result.additional_properties + assert tool.override_called is True + + +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)) + ) + + function_result = await _call_generated_mcp_tool(tool, "slow_op") + + assert function_result.result == "custom fallback summary" + assert function_result.items is not None + 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 function_result.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 @@ -7253,12 +7893,25 @@ async def test_call_tool_as_task_payload_iserror_raises() -> None: side_effect=_send_request_dispatcher( ("tools/call", _make_create_task_result()), ("tasks/get", _make_task_snapshot(status="completed")), - ("tasks/result", _make_payload("payload exploded", is_error=True)), + ( + "tasks/result", + _make_payload( + "payload exploded", + is_error=True, + structured_content={"reason": "task failed"}, + meta={"source": "failed-task"}, + ), + ), ) ) - with pytest.raises(ToolExecutionException, match="payload exploded"): - await tool.call_tool("slow_op") + function_result = await _call_generated_mcp_tool(tool, "slow_op") + + assert function_result.exception is not None + assert function_result.additional_properties["_meta"] == {"source": "failed-task"} + assert function_result.additional_properties[_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY]["structuredContent"] == { + "reason": "task failed" + } async def test_call_tool_as_task_malformed_payload_raises() -> None: