diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index e3ca013f908..79291c7e0db 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -205,6 +205,7 @@ class MCPSpecificApproval(TypedDict, total=False): _MCP_IS_TOOL_KEY = "_mcp_is_tool" _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name" _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY = "_mcp_tool_result_host_payload" +_MCP_TOOL_RESULT_SECURITY_META_KEY = "_mcp_tool_result_security_meta" _MCP_PROGRESSIVE_LIST_TOOL_NAME = "list_mcp_tools" _MCP_PROGRESSIVE_LOAD_TOOL_NAME = "load_tool" _MCP_PROGRESSIVE_UNLOAD_TOOL_NAME = "unload_tool" @@ -251,6 +252,7 @@ class MCPSpecificApproval(TypedDict, total=False): _MCP_SUBSCRIPTION_RETRY_INITIAL_DELAY = 0.1 _MCP_SUBSCRIPTION_RETRY_MAX_DELAY = 2.0 _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES = 1024 * 1024 +_MAX_MCP_SECURITY_META_SIZE_BYTES = 4096 @dataclass @@ -262,8 +264,45 @@ class _MCPHostPayloadCapture: apply_server_meta_to_model_items: bool host_payload: dict[str, Any] | None = None meta: dict[str, Any] | None = None + security_meta: dict[str, Any] | None = None recorded: bool = False meta_prepared: bool = False + security_meta_prepared: bool = False + + def prepare_security_meta(self, mcp_type: Any) -> dict[str, Any] | None: + """Retain security-relevant MCP metadata outside generic payload budgets.""" + if self.security_meta_prepared: + return self.security_meta + self.security_meta_prepared = True + + raw_meta = getattr(mcp_type, "meta", None) + if not isinstance(raw_meta, Mapping) or "ifc" not in raw_meta: + return None + + raw_meta_mapping = cast(Mapping[str, Any], raw_meta) + raw_security_meta: dict[str, Any] = {"ifc": raw_meta_mapping["ifc"]} + if _json_size_exceeds(raw_security_meta, _MAX_MCP_SECURITY_META_SIZE_BYTES): + logger.warning("MCP result _meta.ifc exceeds the independent security metadata limit.") + self.security_meta = {"retained": False} + return self.security_meta + + try: + copied = cast( + dict[str, Any], + to_jsonable_python(raw_security_meta, fallback=str), + ) + except (TypeError, ValueError): + logger.warning("MCP result _meta.ifc could not be retained safely.") + self.security_meta = {"retained": False} + return self.security_meta + + if len(json.dumps(copied).encode("utf-8")) > _MAX_MCP_SECURITY_META_SIZE_BYTES: + logger.warning("MCP result _meta.ifc exceeds the independent security metadata limit.") + self.security_meta = {"retained": False} + return self.security_meta + + self.security_meta = {"retained": True, "meta": copied} + return self.security_meta def prepare_meta(self, mcp_type: Any) -> dict[str, Any] | None: if self.meta_prepared: @@ -293,6 +332,7 @@ def record(self, mcp_type: Any) -> None: if self.recorded: return self.recorded = True + self.prepare_security_meta(mcp_type) self.prepare_meta(mcp_type) effective_limit = self.max_size_bytes @@ -325,8 +365,16 @@ def to_carrier(self) -> _FunctionResultCarrier: return _FunctionResultCarrier( additional_properties=additional_properties, item_additional_properties=item_additional_properties, + private_metadata=( + {_MCP_TOOL_RESULT_SECURITY_META_KEY: self.security_meta} if self.security_meta is not None else {} + ), exclusive_outer_keys=frozenset({_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY}), - exclusive_item_keys=frozenset({"_meta"}), + exclusive_item_keys=frozenset({ + "_meta", + "_security_label_authoritative", + "_security_label_authoritative_confidentiality", + _MCP_TOOL_RESULT_SECURITY_META_KEY, + }), result_already_parsed=True, ) @@ -338,8 +386,11 @@ def prepare_model_result(self, parsed: str | list[Content]) -> list[Content]: updated_item.additional_properties = dict(updated_item.additional_properties) updated_item.additional_properties.pop(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY, None) updated_item.additional_properties.pop("_meta", None) + updated_item.additional_properties.pop(_MCP_TOOL_RESULT_SECURITY_META_KEY, None) if self.apply_server_meta_to_model_items and self.meta is not None: updated_item.additional_properties["_meta"] = self.meta + if self.security_meta is not None: + updated_item.additional_properties[_MCP_TOOL_RESULT_SECURITY_META_KEY] = self.security_meta items[index] = updated_item return items @@ -1413,6 +1464,7 @@ def _parse_tool_result_from_mcp( capture = _mcp_host_payload_capture.get() if capture is not None: + capture.prepare_security_meta(mcp_type) meta = capture.prepare_meta(mcp_type) else: raw_meta = mcp_type.meta diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 0d2e23184d0..5c3dfb2618b 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -22,7 +22,7 @@ Sequence, ) from contextlib import suppress -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone, tzinfo from datetime import time as datetime_time from decimal import Decimal @@ -312,6 +312,10 @@ class _FunctionResultCarrier: additional_properties: dict[str, Any] item_additional_properties: dict[str, Any] + private_metadata: dict[str, Any] = field(default_factory=dict[str, Any]) + parsed_result_postprocessor: Callable[[list[Content]], list[Content]] | None = None + parser_exception_handler: Callable[[], None] | None = None + error_result_finalizer: Callable[[], None] | None = None exclusive_outer_keys: frozenset[str] = frozenset() exclusive_item_keys: frozenset[str] = frozenset() result_already_parsed: bool = False @@ -1218,7 +1222,13 @@ async def invoke( elif callable(configured_parser): try: parsed = configured_parser(result) + if isinstance(parsed, str): + parsed = [Content.from_text(parsed)] + if isinstance(carrier, _FunctionResultCarrier) and carrier.parsed_result_postprocessor is not None: + parsed = carrier.parsed_result_postprocessor(parsed) except Exception as exception: + if isinstance(carrier, _FunctionResultCarrier) and carrier.parser_exception_handler is not None: + carrier.parser_exception_handler() self.invocation_exception_count += 1 if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED: logger.error(f"Function {self.name}: result parser failed. Error: {exception}") @@ -1312,7 +1322,16 @@ async def invoke( elif callable(configured_parser): try: parsed = configured_parser(result) + if isinstance(parsed, str): + parsed = [Content.from_text(parsed)] + if ( + isinstance(carrier, _FunctionResultCarrier) + and carrier.parsed_result_postprocessor is not None + ): + parsed = carrier.parsed_result_postprocessor(parsed) except Exception as exception: + if isinstance(carrier, _FunctionResultCarrier) and carrier.parser_exception_handler is not None: + carrier.parser_exception_handler() self.invocation_exception_count += 1 metric_attributes[OtelAttr.ERROR_TYPE] = type(exception).__name__ if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED: @@ -2164,21 +2183,26 @@ def _finalize_function_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 = merge_additional_properties( - updated_item.additional_properties, - carrier.item_additional_properties, - ) - updated_items[index] = updated_item - function_result.items = updated_items + if carrier is not None and function_result.items is not None: + 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) + is_hidden_variable_reference = ( + updated_item.type == "text" and updated_item.additional_properties.get("_variable_reference") is True + ) + if not is_hidden_variable_reference: + updated_item.additional_properties = merge_additional_properties( + updated_item.additional_properties, + carrier.item_additional_properties, + ) + updated_items[index] = updated_item + function_result.items = updated_items + + if exception is not None and carrier is not None and carrier.error_result_finalizer is not None: + carrier.error_result_finalizer() return function_result diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 83909b1693c..90eeef06303 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -49,11 +49,13 @@ from ._tools import ( _APPROVAL_REQUEST_ID_KEY, # pyright: ignore[reportPrivateUsage] _AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] + _FUNCTION_RESULT_CARRIER_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT, # pyright: ignore[reportPrivateUsage] _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] FunctionTool, _argument_authority_token, # pyright: ignore[reportPrivateUsage] _bind_approval_response_to_pending_request, # pyright: ignore[reportPrivateUsage] + _FunctionResultCarrier, # pyright: ignore[reportPrivateUsage] _load_pending_approval_requests, # pyright: ignore[reportPrivateUsage] tool, ) @@ -119,6 +121,7 @@ _SECURITY_ARGUMENTS_MATERIALIZED_CONTEXT_KEY = "_security_arguments_materialized" _SECURITY_MATERIALIZED_AUTHORITY_CONTEXT_KEY = "_security_materialized_authority" _SECURITY_RESOLVED_AUTHORITY_CONTEXT_KEY = "_security_resolved_authority" +_SECURITY_RESULT_FALLBACK_LABEL_CONTEXT_KEY = "_security_result_fallback_label" _MAX_REPLAY_CONTENT_DEPTH = 32 _REPLAY_STATE_COVERED_KEY = "_agent_framework_fides_replay_state_covered" _REPLAY_STATE_COVERED_MARKER = object() @@ -2116,7 +2119,12 @@ def reset_context_label(self, session: AgentSession | None = None) -> None: ) logger.info("Context label reset to TRUSTED + PUBLIC") - def _update_context_label(self, new_content_label: ContentLabel) -> None: + def _update_context_label( + self, + new_content_label: ContentLabel, + *, + scope: _SecurityScope | None = None, + ) -> None: """Update the context label based on new content added to the context. The context label is updated using the most restrictive policy: @@ -2125,21 +2133,24 @@ def _update_context_label(self, new_content_label: ContentLabel) -> None: Args: new_content_label: The label of the new content being added to context. + scope: Exact security scope to update after task-local activation has unwound. """ - old_label = self._context_label - self._context_label = combine_labels(self._context_label, new_content_label) + target_scope = scope if scope is not None else self._scope + old_label = target_scope.context_label + target_scope.context_label = combine_labels(target_scope.context_label, new_content_label) + updated_label = target_scope.context_label - if old_label != self._context_label: + if old_label != updated_label: logger.info( f">>> CONTEXT TAINT: [{old_label.integrity.value}, {old_label.confidentiality.value}] " - f"-> [{self._context_label.integrity.value}, {self._context_label.confidentiality.value}] " + f"-> [{updated_label.integrity.value}, {updated_label.confidentiality.value}] " f"(new content: [{new_content_label.integrity.value}, {new_content_label.confidentiality.value}])" ) else: logger.debug( "Context label unchanged: [%s, %s]", - self._context_label.integrity.value, - self._context_label.confidentiality.value, + updated_label.integrity.value, + updated_label.confidentiality.value, ) @staticmethod @@ -2755,6 +2766,7 @@ async def process( context_label = self._context_label context.metadata["context_label"] = context_label context.metadata["argument_label"] = argument_label + context.metadata[_SECURITY_RESULT_FALLBACK_LABEL_CONTEXT_KEY] = fallback_label context.metadata["effective_invocation_label"] = combine_labels(context_label, argument_label) materialization_token = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.set(( @@ -5528,18 +5540,18 @@ async def apply_mcp_security_labels( # ``_meta`` keys can be interpreted by additional consumers without touching # the transport layer. _MCP_RESULT_META_KEY = "_meta" +_MCP_RESULT_SECURITY_META_KEY = "_mcp_tool_result_security_meta" def _label_from_mcp_meta(meta: Any) -> ContentLabel | None: """Best-effort parse of an MCP server ``_meta`` payload into a ``ContentLabel``. Looks under ``meta["ifc"]`` for an ``integrity`` and ``confidentiality`` - pair using the wire vocabulary advertised by IFC-aware MCP servers - (e.g. the GitHub MCP server's ``ifc`` schema: integrity is one of - ``"trusted" | "untrusted"`` and confidentiality is one of - ``"public" | "private"``). Returns ``None`` when the payload is missing, - not a dict, or contains values outside the known vocabularies, which lets - callers fall back to a statically derived label. + pair using the wire vocabulary advertised by IFC-aware MCP servers. + ``user_identity`` confidentiality also requires a complete, valid + ``principals`` set. Returns ``None`` when the payload is missing, malformed, + or contains values outside the known vocabularies, which lets callers fail + closed when IFC was supplied. """ if not isinstance(meta, dict): return None @@ -5553,13 +5565,21 @@ def _label_from_mcp_meta(meta: Any) -> ContentLabel | None: try: integrity = IntegrityLabel(integ_raw) if integ_raw is not None else None confidentiality = ConfidentialityLabel(conf_raw) if conf_raw is not None else None - except ValueError: + except (TypeError, ValueError): logger.debug("MCP _meta.ifc had unknown label values: %r", ifc_map) return None if integrity is None or confidentiality is None: logger.debug("MCP _meta.ifc missing integrity/confidentiality: %r", ifc_map) return None - return ContentLabel(integrity=integrity, confidentiality=confidentiality) + metadata: dict[str, Any] | None = None + if confidentiality == ConfidentialityLabel.USER_IDENTITY: + try: + principals = _canonical_principals(ifc_map.get("principals"), source="MCP _meta.ifc") + except ValueError: + logger.debug("MCP _meta.ifc had invalid user-identity principals: %r", ifc_map) + return None + metadata = {PRINCIPAL_METADATA_KEY: _principal_list(principals)} + return ContentLabel(integrity=integrity, confidentiality=confidentiality, metadata=metadata) def _stamp_mcp_content_labels( @@ -5570,13 +5590,13 @@ def _stamp_mcp_content_labels( ) -> Any: """Stamp ``security_label`` on each Content in an MCP tool result. - A parseable server label from ``additional_properties["_meta"]`` can only - restrict ``local_label`` through standard FIDES label combination by - default. When ``trust_server_ifc`` is enabled locally, a complete valid - server label is authoritative for that result. The local label is used - unchanged when metadata is missing, partial, or malformed. The sentinel - ``_meta`` key is consumed regardless so downstream layers do not re-process - it. + A parseable server label can only restrict ``local_label`` through standard + FIDES label combination by default. When ``trust_server_ifc`` is enabled + locally, a complete valid server label is authoritative for that result. + IFC metadata retained by the framework uses an independently bounded + private carrier, so generic Host-payload exhaustion cannot discard it. + Supplied-but-invalid or unretainable IFC metadata fails closed as + UNTRUSTED/PRIVATE. Private carrier keys and ``_meta`` are consumed. """ if not isinstance(contents, list): return contents @@ -5587,9 +5607,35 @@ def _stamp_mcp_content_labels( props = item.additional_properties or {} props.pop(_AUTHORITATIVE_SECURITY_LABEL, None) props.pop(_LEGACY_AUTHORITATIVE_CONFIDENTIALITY, None) + security_meta = props.pop(_MCP_RESULT_SECURITY_META_KEY, None) server_meta = props.pop(_MCP_RESULT_META_KEY, None) - dynamic = _label_from_mcp_meta(server_meta) if server_meta else None - if dynamic is None: + invalid_ifc = False + dynamic: ContentLabel | None = None + if security_meta is not None: + if isinstance(security_meta, Mapping): + security_meta_mapping = cast(Mapping[str, Any], security_meta) + else: + security_meta_mapping = None + if security_meta_mapping is not None and security_meta_mapping.get("retained") is True: + dynamic = _label_from_mcp_meta(security_meta_mapping.get("meta")) + invalid_ifc = dynamic is None + else: + invalid_ifc = True + elif isinstance(server_meta, Mapping) and "ifc" in server_meta: + dynamic = _label_from_mcp_meta(server_meta) + invalid_ifc = dynamic is None + + if invalid_ifc: + label = combine_labels( + local_label, + ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.USER_IDENTITY, + metadata={"source": "invalid_mcp_ifc"}, + ), + ) + props[_AUTHORITATIVE_SECURITY_LABEL] = _INTERNAL_RESULT_MARKER + elif dynamic is None: label = local_label elif trust_server_ifc: label = dynamic @@ -5601,17 +5647,110 @@ def _stamp_mcp_content_labels( return contents_list +def _strip_mcp_authority_markers(contents: str | list[Content]) -> str | list[Content]: + """Remove private authority markers when no tracking middleware can consume them.""" + if not isinstance(contents, list): + return contents + for item in contents: + props = item.additional_properties or {} + props.pop(_AUTHORITATIVE_SECURITY_LABEL, None) + props.pop(_LEGACY_AUTHORITATIVE_CONFIDENTIALITY, None) + item.additional_properties = props + return contents + + +def _stamp_parsed_mcp_content_labels( + context: FunctionInvocationContext, + contents: list[Content], + local_label: ContentLabel, + *, + trust_server_ifc: bool, + retain_authority: bool, +) -> list[Content]: + """Reapply retained MCP IFC after a function-level result parser.""" + raw_carrier = context.metadata.get(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY) + if not isinstance(raw_carrier, _FunctionResultCarrier): + return contents + + security_meta = raw_carrier.private_metadata.get(_MCP_RESULT_SECURITY_META_KEY) + prepared: list[Content] = [] + for item in contents: + updated_item = copy(item) + props = dict(updated_item.additional_properties) + props.pop(_MCP_RESULT_META_KEY, None) + props.pop(_MCP_RESULT_SECURITY_META_KEY, None) + props.pop(_AUTHORITATIVE_SECURITY_LABEL, None) + props.pop(_LEGACY_AUTHORITATIVE_CONFIDENTIALITY, None) + if security_meta is not None: + props[_MCP_RESULT_SECURITY_META_KEY] = security_meta + updated_item.additional_properties = props + prepared.append(updated_item) + + stamped = cast( + list[Content], + _stamp_mcp_content_labels(prepared, local_label, trust_server_ifc=trust_server_ifc), + ) + if not retain_authority: + _strip_mcp_authority_markers(stamped) + return stamped + + +def _stamp_mcp_error_carrier_label( + context: FunctionInvocationContext, + local_label: ContentLabel, + fallback_label: ContentLabel, + *, + trust_server_ifc: bool, +) -> ContentLabel | None: + """Apply retained MCP IFC to a generic function-error result carrier.""" + raw_carrier = context.metadata.get(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY) + if not isinstance(raw_carrier, _FunctionResultCarrier): + return None + + security_meta = raw_carrier.private_metadata.get(_MCP_RESULT_SECURITY_META_KEY) + probe = Content.from_text("") + valid_server_ifc = False + if security_meta is not None: + probe.additional_properties[_MCP_RESULT_SECURITY_META_KEY] = security_meta + if isinstance(security_meta, Mapping): + security_meta_mapping = cast(Mapping[str, Any], security_meta) + if security_meta_mapping.get("retained") is True: + valid_server_ifc = _label_from_mcp_meta(security_meta_mapping.get("meta")) is not None + effective_local_label = ( + fallback_label if security_meta is None or (valid_server_ifc and not trust_server_ifc) else local_label + ) + stamped = cast( + list[Content], + _stamp_mcp_content_labels([probe], effective_local_label, trust_server_ifc=trust_server_ifc), + ) + if not stamped: + return None + + stamped_item = stamped[0] + stamped_properties = dict(stamped_item.additional_properties) + stamped_properties.pop(_AUTHORITATIVE_SECURITY_LABEL, None) + stamped_properties.pop(_LEGACY_AUTHORITATIVE_CONFIDENTIALITY, None) + raw_carrier.item_additional_properties = { + **raw_carrier.item_additional_properties, + **stamped_properties, + } + label_data = stamped_properties.get("security_label") + if not isinstance(label_data, MutableMapping): + return None + return ContentLabel.from_dict(cast(MutableMapping[str, Any], label_data)) + + def _current_mcp_local_label(func_tool: FunctionTool, default_integrity: IntegrityLabel) -> ContentLabel: """Read the current local MCP output policy from a FunctionTool.""" props = func_tool.additional_properties or {} try: integrity = IntegrityLabel(props.get("source_integrity", default_integrity.value)) - except ValueError: - integrity = default_integrity + except (TypeError, ValueError): + integrity = IntegrityLabel.UNTRUSTED try: confidentiality = ConfidentialityLabel(props.get("confidentiality", ConfidentialityLabel.PUBLIC.value)) - except ValueError: - confidentiality = ConfidentialityLabel.PUBLIC + except (TypeError, ValueError): + confidentiality = ConfidentialityLabel.PRIVATE return ContentLabel(integrity=integrity, confidentiality=confidentiality) @@ -5635,10 +5774,72 @@ async def _wrapped(*args: Any, **kwargs: Any) -> Any: local_label = _current_mcp_local_label(func_tool, default_integrity) trust_server_ifc = props.get(_MCP_TRUST_SERVER_IFC_KEY) is True - res = original(*args, **kwargs) - if _inspect.isawaitable(res): - res = await res - return _stamp_mcp_content_labels(res, local_label, trust_server_ifc=trust_server_ifc) + invocation_context = next( + (value for value in kwargs.values() if isinstance(value, FunctionInvocationContext)), + None, + ) + + def prepare_error_label() -> None: + if invocation_context is None: + return + active_context = invocation_context + raw_fallback_label = active_context.metadata.get(_SECURITY_RESULT_FALLBACK_LABEL_CONTEXT_KEY) + fallback_label = raw_fallback_label if isinstance(raw_fallback_label, ContentLabel) else local_label + error_label = _stamp_mcp_error_carrier_label( + active_context, + local_label, + fallback_label, + trust_server_ifc=trust_server_ifc, + ) + active_middleware = _current_middleware.get() + if error_label is None or active_middleware is None or _current_context.get() is not active_context: + return + raw_carrier = active_context.metadata.get(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY) + if not isinstance(raw_carrier, _FunctionResultCarrier): + return + active_scope = active_middleware._scope # pyright: ignore[reportPrivateUsage] + + def commit_error_label() -> None: + active_context.metadata["result_label"] = error_label + active_middleware._update_context_label( # pyright: ignore[reportPrivateUsage] + error_label, + scope=active_scope, + ) + + raw_carrier.error_result_finalizer = commit_error_label + + try: + res = original(*args, **kwargs) + if _inspect.isawaitable(res): + res = await res + except Exception: + prepare_error_label() + raise + + retain_authority = _current_middleware.get() is not None + if invocation_context is not None: + raw_carrier = invocation_context.metadata.get(_FUNCTION_RESULT_CARRIER_CONTEXT_KEY) + if isinstance(raw_carrier, _FunctionResultCarrier): + + def postprocess_parsed_result(contents: list[Content]) -> list[Content]: + return _stamp_parsed_mcp_content_labels( + invocation_context, + contents, + local_label, + trust_server_ifc=trust_server_ifc, + retain_authority=retain_authority, + ) + + raw_carrier.parsed_result_postprocessor = postprocess_parsed_result + raw_carrier.parser_exception_handler = prepare_error_label + + stamped = cast( + str | list[Content], + _stamp_mcp_content_labels(res, local_label, trust_server_ifc=trust_server_ifc), + ) + if not retain_authority: + _strip_mcp_authority_markers(stamped) + return stamped _wrapped._ifc_wrapped = True # type: ignore[attr-defined] func_tool.func = _wrapped diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 0284baeb5ae..e269bd5808d 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -11,8 +11,9 @@ from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping from contextlib import AbstractAsyncContextManager, _AsyncGeneratorContextManager # type: ignore from contextvars import ContextVar +from pathlib import Path from textwrap import dedent -from types import TracebackType +from types import SimpleNamespace, TracebackType from typing import Any, cast from unittest.mock import AsyncMock, Mock, patch @@ -24,17 +25,21 @@ from agent_framework import ( Agent, + AgentSession, ChatResponse, ChatResponseUpdate, Content, + FileSessionStore, FunctionInvocationContext, FunctionMiddleware, FunctionTool, + InMemoryHistoryProvider, MCPStdioTool, MCPStreamableHTTPTool, MCPTaskOptions, # ty: ignore[deprecated] MCPWebsocketTool, # ty: ignore[deprecated] Message, + SessionContext, SupportsChatGetResponse, ) from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning @@ -56,7 +61,7 @@ _should_propagate_cancelled_error, logger, ) -from agent_framework._middleware import FunctionMiddlewarePipeline +from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination from agent_framework._tools import ( _auto_invoke_function, _FunctionResultPayloadBudget, @@ -245,7 +250,10 @@ async def _call_generated_mcp_tool( middleware_pipeline: FunctionMiddlewarePipeline | None = None, host_payload_budget: _FunctionResultPayloadBudget | None = None, mcp_local_label: tuple[str, str] | None = None, + mcp_local_principals: list[dict[str, str]] | None = None, + invocation_session: AgentSession | None = None, trust_server_ifc: bool = False, + include_detailed_errors: bool = False, **kwargs: Any, ) -> Content: function_kwargs: dict[str, Any] = {} @@ -259,7 +267,7 @@ async def _call_generated_mcp_tool( **function_kwargs, ) if mcp_local_label is not None: - from agent_framework.security import IntegrityLabel, _wrap_mcp_function_for_ifc + from agent_framework.security import PRINCIPAL_METADATA_KEY, IntegrityLabel, _wrap_mcp_function_for_ifc source_integrity, confidentiality = mcp_local_label function.additional_properties = { @@ -268,11 +276,16 @@ async def _call_generated_mcp_tool( "confidentiality": confidentiality, "_mcp_trust_server_ifc": trust_server_ifc, } + if mcp_local_principals is not None: + function.additional_properties[PRINCIPAL_METADATA_KEY] = mcp_local_principals _wrap_mcp_function_for_ifc(function, IntegrityLabel(source_integrity)) 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), + config=normalize_function_invocation_configuration({ + "include_detailed_errors": include_detailed_errors, + }), tool_map={tool_name: function}, + invocation_session=invocation_session, middleware_pipeline=middleware_pipeline, host_payload_budget=host_payload_budget, ) @@ -1688,6 +1701,597 @@ def fail_if_copied(*_args: Any, **_kwargs: Any) -> Any: assert _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY not in function_result.additional_properties +async def test_exhausted_aggregate_budget_preserves_ifc_security_label() -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="private result")], + _meta={"ifc": {"integrity": "trusted", "confidentiality": "private"}}, + ) + tool = MCPTool( # type: ignore[abstract] # ty: ignore[call-non-callable] + name="helper", + parse_tool_results=lambda _: "model projection", + max_host_payload_size_bytes=128, + ) + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + budget = _FunctionResultPayloadBudget() + assert budget.reserve(128, 128) + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + host_payload_budget=budget, + mcp_local_label=("trusted", "public"), + trust_server_ifc=True, + ) + + assert "_meta" not in function_result.additional_properties + assert function_result.items is not None + assert function_result.items[0].additional_properties["security_label"] == { + "integrity": "trusted", + "confidentiality": "private", + } + assert "_mcp_tool_result_security_meta" not in function_result.items[0].additional_properties + + +@pytest.mark.parametrize("enable_instrumentation", [False, True], indirect=True) +@pytest.mark.usefixtures("span_exporter") +async def test_function_result_parser_preserves_exhausted_budget_ifc_label() -> None: + from agent_framework.security import LabelTrackingFunctionMiddleware + + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="private result")], + _meta={"ifc": {"integrity": "trusted", "confidentiality": "private"}}, + ) + tool = MCPTool(name="helper", max_host_payload_size_bytes=128) # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + budget = _FunctionResultPayloadBudget() + assert budget.reserve(128, 128) + tracker = LabelTrackingFunctionMiddleware() + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + result_parser=lambda items: items[0].text, + middleware_pipeline=FunctionMiddlewarePipeline(tracker), + host_payload_budget=budget, + mcp_local_label=("trusted", "public"), + trust_server_ifc=True, + ) + + assert function_result.result == "private result" + assert function_result.items is not None + assert function_result.items[0].additional_properties["security_label"] == { + "integrity": "trusted", + "confidentiality": "private", + } + assert tracker.get_context_label().confidentiality.value == "private" + assert "_security_label_authoritative" not in function_result.items[0].additional_properties + json.dumps(function_result.to_dict()) + + +@pytest.mark.parametrize("enable_instrumentation", [False, True], indirect=True) +@pytest.mark.usefixtures("span_exporter") +async def test_function_result_parser_failure_preserves_exhausted_budget_ifc_label() -> None: + from agent_framework.security import LabelTrackingFunctionMiddleware + + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="private parser detail")], + _meta={"ifc": {"integrity": "trusted", "confidentiality": "private"}}, + ) + tool = MCPTool(name="helper", max_host_payload_size_bytes=128) # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + budget = _FunctionResultPayloadBudget() + assert budget.reserve(128, 128) + tracker = LabelTrackingFunctionMiddleware() + + def fail_with_result(items: list[Content]) -> str: + raise ValueError(f"parser rejected {items[0].text}") + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + result_parser=fail_with_result, + middleware_pipeline=FunctionMiddlewarePipeline(tracker), + host_payload_budget=budget, + mcp_local_label=("trusted", "public"), + trust_server_ifc=True, + include_detailed_errors=True, + ) + + assert "private parser detail" in function_result.result + assert function_result.items is not None + assert function_result.items[0].additional_properties["security_label"] == { + "integrity": "trusted", + "confidentiality": "private", + } + assert tracker.get_context_label().confidentiality.value == "private" + assert "_security_label_authoritative" not in function_result.items[0].additional_properties + json.dumps(function_result.to_dict()) + + +@pytest.mark.parametrize( + ("server_ifc", "expected_integrity", "expected_confidentiality"), + [ + pytest.param( + {"integrity": "trusted", "confidentiality": "private"}, + "trusted", + "private", + id="restrictive", + ), + pytest.param( + {"integrity": "trusted"}, + "untrusted", + "user_identity", + id="invalid", + ), + ], +) +async def test_exhausted_aggregate_budget_enforces_ifc_on_detailed_mcp_error( + tmp_path: Path, + server_ifc: dict[str, str], + expected_integrity: str, + expected_confidentiality: str, +) -> None: + from agent_framework.security import ( + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + ) + + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="private MCP error instructions")], + is_error=True, + _meta={"ifc": server_ifc}, + ) + tool = MCPTool( # type: ignore[abstract] # ty: ignore[call-non-callable] + name="helper", + max_host_payload_size_bytes=128, + ) + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + budget = _FunctionResultPayloadBudget() + assert budget.reserve(128, 128) + tracker = LabelTrackingFunctionMiddleware() + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + middleware_pipeline=FunctionMiddlewarePipeline(tracker), + host_payload_budget=budget, + mcp_local_label=("trusted", "public"), + trust_server_ifc=True, + include_detailed_errors=True, + ) + + assert "private MCP error instructions" in function_result.result + assert function_result.items is not None + assert function_result.items[0].additional_properties["security_label"]["integrity"] == expected_integrity + assert ( + function_result.items[0].additional_properties["security_label"]["confidentiality"] == expected_confidentiality + ) + assert "_security_label_authoritative" not in function_result.items[0].additional_properties + assert "_mcp_tool_result_security_meta" not in function_result.items[0].additional_properties + json.dumps(function_result.to_dict()) + store = FileSessionStore(tmp_path / expected_integrity, serialization_format="json") + session = AgentSession(session_id=f"mcp-error-{expected_integrity}") + session.state["mcp_error_result"] = function_result + await store.set(session.session_id, session) + restored = await store.get(session.session_id) + assert restored is not None + restored_result = restored.state["mcp_error_result"] + assert isinstance(restored_result, dict) + restored_items = restored_result["items"] + assert isinstance(restored_items, list) + assert restored_items[0]["additional_properties"]["security_label"]["integrity"] == expected_integrity + context_label = tracker.get_context_label() + assert context_label.integrity.value == expected_integrity + assert context_label.confidentiality.value == expected_confidentiality + + protected_tool = FunctionTool( + func=lambda: "must not run", + name="protected_mcp_error_sink", + input_model={"type": "object", "properties": {}}, + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) + invocation = FunctionInvocationContext(function=protected_tool, arguments={}) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, PolicyEnforcementFunctionMiddleware()).execute( + invocation, + lambda _context: pytest.fail("MCP error IFC must block the protected public sink"), + ) + + +async def test_retried_mcp_error_does_not_commit_swallowed_label() -> None: + from agent_framework.security import ( + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + ) + + error_result = types.CallToolResult( + content=[types.TextContent(type="text", text="private failed attempt")], + is_error=True, + _meta={"ifc": {"integrity": "trusted", "confidentiality": "private"}}, + ) + success_result = types.CallToolResult( + content=[types.TextContent(type="text", text="recovered")], + is_error=False, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(side_effect=[error_result, success_result]) + + class RetryOnceMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Any]) -> None: + del context + try: + await call_next() + except ToolExecutionException: + await call_next() + + tracker = LabelTrackingFunctionMiddleware() + function_result = await _call_generated_mcp_tool( + tool, + "widget", + middleware_pipeline=FunctionMiddlewarePipeline(RetryOnceMiddleware(), tracker), + mcp_local_label=("trusted", "public"), + trust_server_ifc=True, + ) + + assert function_result.result == "recovered" + context_label = tracker.get_context_label() + assert context_label.integrity.value == "trusted" + assert context_label.confidentiality.value == "public" + + executed = False + + def public_sink() -> str: + nonlocal executed + executed = True + return "sent" + + sink_tool = FunctionTool( + func=public_sink, + name="public_after_retry_sink", + input_model={"type": "object", "properties": {}}, + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) + invocation = FunctionInvocationContext(function=sink_tool, arguments={}) + + async def execute_public_sink(_: FunctionInvocationContext) -> str: + return public_sink() + + await FunctionMiddlewarePipeline(tracker, PolicyEnforcementFunctionMiddleware()).execute( + invocation, + execute_public_sink, + ) + assert executed is True + + +async def test_mcp_error_finalizer_updates_invocation_session_scope() -> None: + from agent_framework.security import ( + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + ) + + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="session-private failure")], + is_error=True, + _meta={"ifc": {"integrity": "trusted", "confidentiality": "private"}}, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + tracker = LabelTrackingFunctionMiddleware() + session = AgentSession(session_id="mcp-error-session-scope") + + await _call_generated_mcp_tool( + tool, + "widget", + middleware_pipeline=FunctionMiddlewarePipeline(tracker), + mcp_local_label=("trusted", "public"), + invocation_session=session, + trust_server_ifc=True, + include_detailed_errors=True, + ) + + session_label = tracker.get_context_label(session) + assert session_label.integrity.value == "trusted" + assert session_label.confidentiality.value == "private" + default_label = tracker.get_context_label() + assert default_label.integrity.value == "trusted" + assert default_label.confidentiality.value == "public" + + protected_tool = FunctionTool( + func=lambda: "must not run", + name="session_public_after_mcp_error_sink", + input_model={"type": "object", "properties": {}}, + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) + invocation = FunctionInvocationContext(function=protected_tool, arguments={}, session=session) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, PolicyEnforcementFunctionMiddleware()).execute( + invocation, + lambda _context: pytest.fail("The invocation session must retain the MCP error restriction"), + ) + + +async def test_mcp_error_without_ifc_preserves_local_identity_principals() -> None: + from agent_framework.security import ( + PRINCIPAL_METADATA_KEY, + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + ) + + principals = [{"tenant_id": "tenant-a", "user_id": "user-a"}] + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="identity-scoped failure")], + is_error=True, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + tracker = LabelTrackingFunctionMiddleware() + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + middleware_pipeline=FunctionMiddlewarePipeline(tracker), + mcp_local_label=("trusted", "user_identity"), + mcp_local_principals=principals, + include_detailed_errors=True, + ) + + assert function_result.items is not None + result_label = function_result.items[0].additional_properties["security_label"] + assert result_label["integrity"] == "trusted" + assert result_label["confidentiality"] == "user_identity" + assert result_label["metadata"][PRINCIPAL_METADATA_KEY] == principals + context_label = tracker.get_context_label() + assert context_label.metadata[PRINCIPAL_METADATA_KEY] == principals + + executed = False + + def identity_sink() -> str: + nonlocal executed + executed = True + return "sent" + + sink_tool = FunctionTool( + func=identity_sink, + name="identity_after_mcp_error_sink", + input_model={"type": "object", "properties": {}}, + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "user_identity", + PRINCIPAL_METADATA_KEY: principals, + }, + ) + invocation = FunctionInvocationContext(function=sink_tool, arguments={}) + + async def execute_identity_sink(_: FunctionInvocationContext) -> str: + return identity_sink() + + await FunctionMiddlewarePipeline(tracker, PolicyEnforcementFunctionMiddleware()).execute( + invocation, + execute_identity_sink, + ) + assert executed is True + + +async def test_restriction_only_mcp_error_preserves_private_argument_fallback() -> None: + from agent_framework.security import ( + ConfidentialityLabel, + ContentLabel, + IntegrityLabel, + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + ) + + secret = "private argument echoed by MCP error" + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text=secret)], + is_error=True, + _meta={"ifc": {"integrity": "trusted", "confidentiality": "public"}}, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + tracker = LabelTrackingFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + secret, + ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ), + ) + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + middleware_pipeline=FunctionMiddlewarePipeline(tracker), + mcp_local_label=("trusted", "public"), + include_detailed_errors=True, + value=f"[{variable_id}]", + ) + + assert secret in function_result.result + assert function_result.items is not None + result_label = function_result.items[0].additional_properties["security_label"] + assert result_label["integrity"] == "trusted" + assert result_label["confidentiality"] == "private" + context_label = tracker.get_context_label() + assert context_label.integrity == IntegrityLabel.TRUSTED + assert context_label.confidentiality == ConfidentialityLabel.PRIVATE + + protected_tool = FunctionTool( + func=lambda: "must not run", + name="public_after_private_mcp_error_sink", + input_model={"type": "object", "properties": {}}, + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) + invocation = FunctionInvocationContext(function=protected_tool, arguments={}) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, PolicyEnforcementFunctionMiddleware()).execute( + invocation, + lambda _context: pytest.fail("Private MCP error data must block the public sink"), + ) + + +async def test_oversized_identity_ifc_fails_closed_without_confidentiality_downgrade() -> None: + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="identity-scoped result")], + _meta={ + "ifc": { + "integrity": "trusted", + "confidentiality": "user_identity", + "principals": [{"tenant_id": "tenant-a", "user_id": "x" * 5000}], + } + }, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + + function_result = await _call_generated_mcp_tool( + tool, + "widget", + host_payload_budget=_FunctionResultPayloadBudget(), + mcp_local_label=("trusted", "public"), + trust_server_ifc=True, + ) + + assert function_result.items is not None + assert function_result.items[0].additional_properties["security_label"] == { + "integrity": "untrusted", + "confidentiality": "user_identity", + "metadata": {"source": "invalid_mcp_ifc"}, + } + assert "_security_label_authoritative" not in function_result.items[0].additional_properties + assert "_mcp_tool_result_security_meta" not in function_result.items[0].additional_properties + json.dumps(function_result.to_dict()) + + +@pytest.mark.parametrize( + "server_principals", + [ + pytest.param(None, id="missing"), + pytest.param([], id="empty"), + pytest.param([{"tenant_id": "tenant-a"}], id="missing-user"), + pytest.param([{"tenant_id": "", "user_id": "user-a"}], id="empty-tenant"), + pytest.param("tenant-a/user-a", id="not-a-sequence-of-mappings"), + ], +) +async def test_invalid_identity_ifc_cannot_inherit_local_principals(server_principals: Any) -> None: + from agent_framework.security import ( + PRINCIPAL_METADATA_KEY, + ConfidentialityLabel, + IntegrityLabel, + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + _wrap_mcp_function_for_ifc, + ) + + principals = [{"tenant_id": "tenant-a", "user_id": "user-a"}] + server_ifc: dict[str, Any] = { + "integrity": "trusted", + "confidentiality": "user_identity", + } + if server_principals is not None: + server_ifc["principals"] = server_principals + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="identity-scoped result")], + _meta={"ifc": server_ifc}, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_session.call_tool = AsyncMock(return_value=mcp_result) + function = FunctionTool( + name="identity_widget", + description="", + func=_make_mcp_tool_caller(tool, "identity_widget"), + input_model={"type": "object", "properties": {}}, + additional_properties={ + "_mcp_remote_name": "identity_widget", + "source_integrity": "trusted", + "confidentiality": "user_identity", + PRINCIPAL_METADATA_KEY: principals, + }, + ) + _wrap_mcp_function_for_ifc(function, IntegrityLabel.TRUSTED) + tracker = LabelTrackingFunctionMiddleware(auto_hide_untrusted=True) + policy = PolicyEnforcementFunctionMiddleware() + + result = await _auto_invoke_function( + Content.from_function_call(call_id="identity-call", name=function.name, arguments={}), + config=normalize_function_invocation_configuration(None), + tool_map={function.name: function}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + host_payload_budget=_FunctionResultPayloadBudget(), + ) + + assert result.items is not None + assert result.items[0].additional_properties["_variable_reference"] is True + context_label = tracker.get_context_label() + assert context_label.confidentiality == ConfidentialityLabel.USER_IDENTITY + assert PRINCIPAL_METADATA_KEY not in context_label.metadata + + executed = False + + def identity_destination() -> str: + nonlocal executed + executed = True + return "sent" + + destination = FunctionTool( + name="identity_destination", + description="Identity-scoped destination", + func=identity_destination, + input_model={"type": "object", "properties": {}}, + additional_properties={ + "accepts_untrusted": True, + "source_integrity": "trusted", + "max_allowed_confidentiality": "user_identity", + PRINCIPAL_METADATA_KEY: principals, + }, + ) + + with pytest.raises(MiddlewareTermination): + await _auto_invoke_function( + Content.from_function_call(call_id="destination-call", name=destination.name, arguments={}), + config=normalize_function_invocation_configuration(None), + tool_map={destination.name: destination}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) + + assert executed is False + + 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")], @@ -1978,10 +2582,122 @@ async def test_secure_mcp_builtin_parser_restricts_all_result_shapes(result_shap assert hidden_item.additional_properties["_variable_reference"] is True assert hidden_item.additional_properties["security_label"]["integrity"] == "untrusted" assert hidden_item.additional_properties["security_label"]["confidentiality"] == "private" - assert hidden_item.additional_properties["_meta"] == mcp_result.meta + assert "_meta" not in hidden_item.additional_properties assert function_result.additional_properties["_meta"] == mcp_result.meta +async def test_hidden_builtin_mcp_result_replays_without_tainting_context() -> None: + """A restored hidden MCP placeholder stays integrity-clean on the next run.""" + from agent_framework.security import ( + IntegrityLabel, + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + SecureAgentConfig, + _wrap_mcp_function_for_ifc, + ) + + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="untrusted MCP payload")], + _meta={"ifc": {"integrity": "untrusted", "confidentiality": "private"}}, + ) + tool = MCPTool(name="helper") # type: ignore[abstract] # ty: ignore[call-non-callable] + mock_session = Mock() + tool.session = mock_session + mock_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", + "confidentiality": "private", + "max_allowed_confidentiality": "public", + }, + ) + _wrap_mcp_function_for_ifc(function, IntegrityLabel.UNTRUSTED) + security = SecureAgentConfig() + session = AgentSession(session_id="hidden-mcp-replay") + first_context = SessionContext( + session_id=session.session_id, + input_messages=[Message(role="user", contents=["run"])], + ) + await security.before_run( + agent=cast(Any, SimpleNamespace(context_providers=[security])), + session=session, + context=first_context, + state=session.state.setdefault(security.source_id, {}), + ) + first_tracker = first_context.get_middleware()[0] + assert isinstance(first_tracker, LabelTrackingFunctionMiddleware) + + 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}, + invocation_session=session, + middleware_pipeline=FunctionMiddlewarePipeline(first_tracker), + host_payload_budget=_FunctionResultPayloadBudget(), + ) + + assert function_result.items is not None + assert function_result.items[0].additional_properties["_variable_reference"] is True + assert "_meta" not in function_result.items[0].additional_properties + assert function_result.additional_properties["_meta"] == mcp_result.meta + + restored = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + history = InMemoryHistoryProvider() + replayed = Message(role="tool", contents=[function_result]) + second_context = SessionContext(session_id=restored.session_id, input_messages=[]) + second_context.extend_messages(history.source_id, [replayed]) + await security.before_run( + agent=cast(Any, SimpleNamespace(context_providers=[history, security])), + session=restored, + context=second_context, + state=restored.state.setdefault(security.source_id, {}), + ) + second_tracker, second_policy, _ = second_context.get_middleware() + assert isinstance(second_tracker, LabelTrackingFunctionMiddleware) + assert isinstance(second_policy, PolicyEnforcementFunctionMiddleware) + restored_label = second_tracker.get_context_label() + assert restored_label.integrity.value == "trusted" + assert restored_label.confidentiality.value == "private" + + executed = False + + def protected_action() -> str: + nonlocal executed + executed = True + return "done" + + protected_tool = FunctionTool( + name="protected_action", + description="Run an unrelated protected action", + func=protected_action, + input_model={"type": "object", "properties": {}}, + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "private", + }, + ) + invocation = FunctionInvocationContext( + function=protected_tool, + arguments={}, + session=restored, + ) + + async def execute_protected(_: FunctionInvocationContext) -> list[Content]: + return [Content.from_text(protected_action())] + + await FunctionMiddlewarePipeline(second_tracker, second_policy).execute( + invocation, + execute_protected, + ) + + assert executed is True + + async def test_secure_mcp_builtin_parser_honors_locally_trusted_server_ifc() -> None: from agent_framework.security import LabelTrackingFunctionMiddleware diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 6c7ec046fe6..018c0b4d15d 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -8096,6 +8096,22 @@ def test_label_from_meta_trusted_public(self): assert label.integrity == IntegrityLabel.TRUSTED assert label.confidentiality == ConfidentialityLabel.PUBLIC + def test_label_from_meta_user_identity_requires_and_preserves_principals(self) -> None: + from agent_framework.security import _label_from_mcp_meta + + principals = [{"tenant_id": "tenant-a", "user_id": "user-a"}] + label = _label_from_mcp_meta({ + "ifc": { + "integrity": "trusted", + "confidentiality": "user_identity", + "principals": principals, + } + }) + + assert label is not None + assert label.confidentiality == ConfidentialityLabel.USER_IDENTITY + assert label.metadata[_PRINCIPALS_KEY] == principals + @pytest.mark.parametrize( "meta", [ @@ -8107,6 +8123,15 @@ def test_label_from_meta_trusted_public(self): {"ifc": {"integrity": "trusted"}}, # missing confidentiality {"ifc": {"integrity": "garbage", "confidentiality": "public"}}, {"ifc": {"integrity": "trusted", "confidentiality": "garbage"}}, + {"ifc": {"integrity": "trusted", "confidentiality": "user_identity"}}, + {"ifc": {"integrity": "trusted", "confidentiality": "user_identity", "principals": []}}, + { + "ifc": { + "integrity": "trusted", + "confidentiality": "user_identity", + "principals": [{"tenant_id": "tenant-a"}], + } + }, {"other_key": {"integrity": "trusted", "confidentiality": "public"}}, # non-ifc key ], ) @@ -8174,36 +8199,45 @@ def test_stamp_contents_complete_local_and_remote_label_matrix(self, trust_serve ConfidentialityLabel.PRIVATE: 1, ConfidentialityLabel.USER_IDENTITY: 2, } + local_principals = [{"tenant_id": "tenant-local", "user_id": "user-local"}] + remote_principals = [{"tenant_id": "tenant-remote", "user_id": "user-remote"}] for local_integrity in integrity_labels: for local_confidentiality in confidentiality_labels: for remote_integrity in integrity_labels: for remote_confidentiality in confidentiality_labels: + local_metadata: dict[str, Any] = {"source": "local_mcp_policy"} + if local_confidentiality == ConfidentialityLabel.USER_IDENTITY: + local_metadata[_PRINCIPALS_KEY] = local_principals static = ContentLabel( integrity=local_integrity, confidentiality=local_confidentiality, - metadata={"source": "local_mcp_policy"}, + metadata=local_metadata, ) + remote_ifc: dict[str, Any] = { + "integrity": remote_integrity.value, + "confidentiality": remote_confidentiality.value, + "metadata": {"source": "forged_remote_policy"}, + } + if remote_confidentiality == ConfidentialityLabel.USER_IDENTITY: + remote_ifc["principals"] = remote_principals contents = [ Content.from_text( "x", - additional_properties={ - "_meta": { - "ifc": { - "integrity": remote_integrity.value, - "confidentiality": remote_confidentiality.value, - "metadata": {"source": "forged_remote_policy"}, - } - } - }, + additional_properties={"_meta": {"ifc": remote_ifc}}, ) ] _stamp_mcp_content_labels(contents, static, trust_server_ifc=trust_server_ifc) + expected_metadata: dict[str, Any] | None if trust_server_ifc: expected_integrity = remote_integrity expected_confidentiality = remote_confidentiality - expected_metadata: dict[str, Any] | None = None + expected_metadata = ( + {_PRINCIPALS_KEY: remote_principals} + if remote_confidentiality == ConfidentialityLabel.USER_IDENTITY + else None + ) else: expected_integrity = ( IntegrityLabel.UNTRUSTED @@ -8214,6 +8248,19 @@ def test_stamp_contents_complete_local_and_remote_label_matrix(self, trust_serve (local_confidentiality, remote_confidentiality), key=confidentiality_rank.__getitem__ ) expected_metadata = {"source": "local_mcp_policy"} + if expected_confidentiality == ConfidentialityLabel.USER_IDENTITY: + expected_metadata[_PRINCIPALS_KEY] = [ + *( + local_principals + if local_confidentiality == ConfidentialityLabel.USER_IDENTITY + else [] + ), + *( + remote_principals + if remote_confidentiality == ConfidentialityLabel.USER_IDENTITY + else [] + ), + ] assert contents[0].additional_properties["security_label"] == { "integrity": expected_integrity.value, "confidentiality": expected_confidentiality.value, @@ -8248,9 +8295,7 @@ def test_stamp_contents_missing_meta_falls_back_to_local_policy( ids=["malformed", "partial"], ) @pytest.mark.parametrize("trust_server_ifc", [False, True]) - def test_stamp_contents_invalid_meta_falls_back_to_local_policy( - self, server_meta: dict[str, Any], trust_server_ifc: bool - ): + def test_stamp_contents_invalid_meta_fails_closed(self, server_meta: dict[str, Any], trust_server_ifc: bool): from agent_framework.security import _stamp_mcp_content_labels static = ContentLabel(integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC) @@ -8262,11 +8307,67 @@ def test_stamp_contents_invalid_meta_falls_back_to_local_policy( ] _stamp_mcp_content_labels(contents, static, trust_server_ifc=trust_server_ifc) assert contents[0].additional_properties["security_label"] == { - "integrity": "trusted", - "confidentiality": "public", + "integrity": "untrusted", + "confidentiality": "user_identity", + "metadata": {"source": "invalid_mcp_ifc"}, } assert "_meta" not in contents[0].additional_properties + def test_stamp_contents_unretained_ifc_fails_closed(self) -> None: + """A framework-dropped IFC label cannot silently fall back to public.""" + from agent_framework.security import _stamp_mcp_content_labels + + static = ContentLabel(integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC) + contents = [ + Content.from_text( + "x", + additional_properties={"_mcp_tool_result_security_meta": {"retained": False}}, + ) + ] + + _stamp_mcp_content_labels(contents, static, trust_server_ifc=True) + + assert contents[0].additional_properties["security_label"] == { + "integrity": "untrusted", + "confidentiality": "user_identity", + "metadata": {"source": "invalid_mcp_ifc"}, + } + assert "_mcp_tool_result_security_meta" not in contents[0].additional_properties + + async def test_wrap_mcp_function_invalid_meta_is_serializable_without_tracking(self) -> None: + from agent_framework.security import _wrap_mcp_function_for_ifc + + async def fake_call(**kwargs: Any) -> list[Content]: + return [ + Content.from_text( + "payload", + additional_properties={"_meta": {"ifc": {"integrity": "trusted"}}}, + ) + ] + + function = FunctionTool( + func=fake_call, + name="remote_tool", + description="", + additional_properties={ + "source_integrity": "trusted", + "confidentiality": "public", + "_mcp_remote_name": "remote_tool", + }, + ) + _wrap_mcp_function_for_ifc(function, IntegrityLabel.TRUSTED) + assert function.func is not None + + result = await function.func() + + assert result[0].additional_properties["security_label"] == { + "integrity": "untrusted", + "confidentiality": "user_identity", + "metadata": {"source": "invalid_mcp_ifc"}, + } + assert _AUTHORITY_MARKER_KEY not in result[0].additional_properties + json.dumps(result[0].to_dict()) + def test_stamp_contents_non_ifc_meta_falls_back_to_static(self): """Generic ``_meta`` keys unrelated to IFC don't accidentally produce a label.""" from agent_framework.security import _stamp_mcp_content_labels @@ -8599,13 +8700,13 @@ async def test_framework_stamped_mcp_label_remains_authoritative_through_trackin server_meta=server_meta, ) await apply_mcp_security_labels(mcp_tool, trust_server_ifc=True) - assert function.func is not None - stamped_result = await function.func() + wrapped = function.func + assert wrapped is not None tracker = LabelTrackingFunctionMiddleware() context = FunctionInvocationContext(function=function, arguments={}) async def next_fn() -> None: - context.result = stamped_result + context.result = await wrapped() await tracker.process(context, next_fn)