diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 96478b9eb4..f546cdda43 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -2,6 +2,8 @@ from __future__ import annotations +import asyncio +import contextvars import inspect import logging import sys @@ -38,6 +40,10 @@ logger = logging.getLogger(__name__) +_ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT: contextvars.ContextVar[tuple[Any, asyncio.Task[Any] | None] | None] = ( + contextvars.ContextVar("agent_framework_core_active_function_middleware_context", default=None) +) + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -1570,9 +1576,20 @@ def create_next_handler(index: int) -> Callable[[], Awaitable[None]]: if index >= len(self._middleware): async def final_wrapper() -> None: - context.result = final_handler(context) - if inspect.isawaitable(context.result): - context.result = await context.result + from ._tools import ( + _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT, # pyright: ignore[reportPrivateUsage] + _suspend_security_argument_materialization, # pyright: ignore[reportPrivateUsage] + ) + + security_materialization = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.get() + if security_materialization is not None: + active_context, materialize, origin_task = security_materialization + if context is active_context and asyncio.current_task() is origin_task: + materialize(context) + with _suspend_security_argument_materialization(): + context.result = final_handler(context) + if inspect.isawaitable(context.result): + context.result = await context.result return final_wrapper @@ -1584,7 +1601,11 @@ async def current_handler() -> None: first_handler = create_next_handler(0) # Don't suppress MiddlewareTermination - let it propagate to signal loop termination - await first_handler() + pipeline_token = _ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT.set((context, asyncio.current_task())) + try: + await first_handler() + finally: + _ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT.reset(pipeline_token) return context.result diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 8aa177ec11..126d3307a9 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -17,11 +17,12 @@ AsyncIterable, Awaitable, Callable, + Generator, Iterable, Mapping, Sequence, ) -from contextlib import suppress +from contextlib import contextmanager, suppress from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone, tzinfo from datetime import time as datetime_time @@ -138,12 +139,25 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None) _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY: Final[str] = "_security_function_arguments" _PREPARED_ARGUMENTS_CONTEXT_KEY: Final[str] = "_prepared_function_arguments" _AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY: Final[str] = "_auto_prepare_function_arguments" +_SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT: Final[ + contextvars.ContextVar[tuple[Any, Callable[[Any], None], asyncio.Task[Any] | None] | None] +] = contextvars.ContextVar("agent_framework_core_security_argument_materialization_context", default=None) _FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT: Final[str] = ( "Function invocation limit reached before a final answer could be produced." ) _USER_VISIBLE_CONTENT_TYPES: Final[set[str]] = {"data", "uri", "error", "hosted_file", "hosted_vector_store"} +@contextmanager +def _suspend_security_argument_materialization() -> Generator[None]: + """Prevent nested tool calls from inheriting one invocation's materialization capability.""" + token = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.set(None) + try: + yield + finally: + _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.reset(token) + + class _FunctionArgumentValidationError(TypeError): """An argument-validation failure raised before the function body starts.""" @@ -1167,7 +1181,6 @@ async def invoke( direct_argument_kwargs = ( {key: value for key, value in kwargs.items() if key in parameter_names} if arguments is None else {} ) - runtime_kwargs = dict(context.kwargs) if context is not None else {} unexpected_kwargs = {key: value for key, value in kwargs.items() if key not in direct_argument_kwargs} if unexpected_kwargs: unexpected_names = ", ".join(sorted(unexpected_kwargs)) @@ -1180,6 +1193,23 @@ async def invoke( if arguments is None and context is not None: arguments = context.arguments + security_materialization = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.get() + if security_materialization is not None: + active_context, materialize, origin_task = security_materialization + if ( + asyncio.current_task() is origin_task + and active_context.function is self + and (context is None or context is active_context) + ): + direct_arguments = self._arguments_as_mapping(arguments) + if direct_arguments is not None: + active_context.arguments = direct_arguments + active_context.metadata[_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY] = True + materialize(active_context) + arguments = active_context.arguments + context = active_context + + runtime_kwargs = dict(context.kwargs) if context is not None else {} current_arguments = self._arguments_as_mapping(arguments) approval_visible_arguments = self._approval_visible_arguments(arguments, context) self._ensure_security_arguments_unchanged(context, current_arguments) @@ -1211,7 +1241,8 @@ async def invoke( if not OBSERVABILITY_SETTINGS.ENABLED: logger.info(f"Function name: {self.name}") logger.debug(f"Function arguments: {observable_kwargs}") - result = await self._invoke_function(call_kwargs) + with _suspend_security_argument_materialization(): + result = await self._invoke_function(call_kwargs) if skip_parsing: logger.info(f"Function {self.name} succeeded.") logger.debug(f"Function result: {type(result).__name__}") @@ -1290,7 +1321,8 @@ async def invoke( start_time_stamp = perf_counter() end_time_stamp: float | None = None try: - result = await self._invoke_function(call_kwargs) + with _suspend_security_argument_materialization(): + result = await self._invoke_function(call_kwargs) end_time_stamp = perf_counter() except Exception as exception: end_time_stamp = perf_counter() diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 954b7ce304..931c70d9e4 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -28,6 +28,7 @@ from copy import copy, deepcopy from dataclasses import dataclass, replace from datetime import datetime, timedelta +from difflib import get_close_matches from enum import Enum from typing import TYPE_CHECKING, Annotated, Any, NoReturn, cast from weakref import WeakKeyDictionary @@ -35,12 +36,19 @@ from pydantic import BaseModel, Field from ._feature_stage import ExperimentalFeature, experimental -from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareFailure, MiddlewareTermination +from ._middleware import ( + _ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT, # pyright: ignore[reportPrivateUsage] + FunctionInvocationContext, + FunctionMiddleware, + MiddlewareFailure, + MiddlewareTermination, +) from ._serialization import SerializationMixin from ._sessions import AgentSession, ContextProvider from ._tools import ( _APPROVAL_REQUEST_ID_KEY, # pyright: ignore[reportPrivateUsage] _AUTO_ARGUMENT_PREPARATION_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] @@ -104,6 +112,23 @@ # with the content and break the lookup. _VARIABLE_ID_CONSUMERS = frozenset({"inspect_variable", "quarantined_llm"}) _REWRITTEN_ARGUMENT_INDICES_KEY = "_rewritten_argument_indices" +_SECURITY_POLICY_ERRORS_CONTEXT_KEY = "_agent_framework_core_security_policy_errors" +_SECURITY_ARGUMENTS_MATERIALIZED_CONTEXT_KEY = "_agent_framework_core_security_arguments_materialized" +_SECURITY_ARGUMENTS_REQUIRE_MATERIALIZATION_CONTEXT_KEY = ( + "_agent_framework_core_security_arguments_require_materialization" +) +_SECURITY_MATERIALIZED_AUTHORITY_CONTEXT_KEY = "_agent_framework_core_security_materialized_authority" +_SECURITY_RESOLVED_AUTHORITY_CONTEXT_KEY = "_agent_framework_core_security_resolved_authority" +_TOOL_POLICY_ACCEPTS_UNTRUSTED_KEY = "accepts_untrusted" +_TOOL_POLICY_CONFIDENTIALITY_KEY = "confidentiality" +_TOOL_POLICY_MAX_ALLOWED_CONFIDENTIALITY_KEY = "max_allowed_confidentiality" +_TOOL_POLICY_SOURCE_INTEGRITY_KEY = "source_integrity" +_TOOL_SECURITY_POLICY_KEYS = frozenset({ + _TOOL_POLICY_ACCEPTS_UNTRUSTED_KEY, + _TOOL_POLICY_CONFIDENTIALITY_KEY, + _TOOL_POLICY_MAX_ALLOWED_CONFIDENTIALITY_KEY, + _TOOL_POLICY_SOURCE_INTEGRITY_KEY, +}) def _get_additional_properties(obj: Any) -> dict[str, Any]: @@ -112,32 +137,35 @@ def _get_additional_properties(obj: Any) -> dict[str, Any]: return cast(dict[str, Any], props) if isinstance(props, dict) else {} -def _top_level_argument_value(context: FunctionInvocationContext, arg_name: str) -> tuple[Any, str | None]: - """Locate a top-level argument value in either context.arguments or context.kwargs. +def _top_level_argument_value( + arguments: Any, + kwargs: Mapping[str, Any], + arg_name: str, +) -> tuple[Any, str | None]: + """Locate a top-level argument value in either arguments or runtime kwargs. Returns the value and a string indicating its source ('arguments' or 'kwargs'), or (None, None) if not found. """ - args = cast(Any, context.arguments) - if isinstance(args, Mapping) and arg_name in args: - return cast(Any, args[arg_name]), "arguments" + args = cast(object, arguments) + if isinstance(args, Mapping): + args_mapping = cast(Mapping[str, Any], args) + if arg_name in args_mapping: + return args_mapping[arg_name], "arguments" - kwargs = cast(Any, context.kwargs) - if isinstance(kwargs, Mapping) and arg_name in kwargs: - return cast(Any, kwargs[arg_name]), "kwargs" + if arg_name in kwargs: + return kwargs[arg_name], "kwargs" return None, None -def _top_level_argument_keys(context: FunctionInvocationContext) -> set[str]: - """Union of top level keys in context.arguments and context.kwargs.""" +def _top_level_argument_keys(arguments: Any, kwargs: Mapping[str, Any]) -> set[str]: + """Return the union of top-level argument and runtime-kwarg keys.""" keys: set[str] = set() - args = cast(Any, context.arguments) + args = cast(object, arguments) if isinstance(args, Mapping): keys.update(cast(Mapping[str, Any], args).keys()) - kwargs = cast(Any, context.kwargs) - if isinstance(kwargs, Mapping): - keys.update(cast(Mapping[str, Any], kwargs).keys()) + keys.update(kwargs.keys()) return keys @@ -417,6 +445,113 @@ def combine_labels(*labels: ContentLabel) -> ContentLabel: ) +@dataclass(frozen=True, slots=True) +class _ToolSecurityPolicy: + """Parsed FIDES tool metadata used consistently by tracking and enforcement.""" + + source_integrity: IntegrityLabel | None + confidentiality: ConfidentialityLabel + max_allowed_confidentiality: ConfidentialityLabel | None + accepts_untrusted: bool + principal_state: str + principals: tuple[_Principal, ...] | None + internal_security_tool: bool + errors: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _SecurityInvocationAuthority: + """Exact security authority inspected for one function invocation.""" + + function: FunctionTool + tool_policy: _ToolSecurityPolicy + session: AgentSession | None + security_scope: _SecurityScope + arguments_token: Any + + +@dataclass(frozen=True, slots=True) +class _ResolvedInvocationArguments: + """Detached variable expansion used for policy preflight and later materialization.""" + + arguments: Any + kwargs: dict[str, Any] + labels: tuple[ContentLabel, ...] + rewritten_arguments: dict[str, set[int]] + authority_token: _SecurityInvocationAuthority + + +def _parse_tool_security_policy( + function: FunctionTool, + *, + default_confidentiality: ConfidentialityLabel = ConfidentialityLabel.PUBLIC, +) -> _ToolSecurityPolicy: + """Parse FIDES metadata without silently accepting unknown values.""" + props = _get_additional_properties(function) + errors: list[str] = [] + for key in props: + if not isinstance(key, str) or key in _TOOL_SECURITY_POLICY_KEYS: + continue + matches = get_close_matches(key.casefold(), sorted(_TOOL_SECURITY_POLICY_KEYS), n=1, cutoff=0.9) + if matches: + errors.append(f"unknown FIDES policy key {key!r}; did you mean {matches[0]!r}?") + + def parse_enum( + key: str, + enum_type: type[IntegrityLabel] | type[ConfidentialityLabel], + ) -> IntegrityLabel | ConfidentialityLabel | None: + if key not in props: + return None + raw_value = props[key] + if not isinstance(raw_value, str): + errors.append(f"{key} must be a string") + return None + try: + return enum_type(raw_value) + except ValueError: + errors.append(f"{key} has unknown value {raw_value!r}") + return None + + source_integrity = parse_enum(_TOOL_POLICY_SOURCE_INTEGRITY_KEY, IntegrityLabel) + confidentiality = parse_enum(_TOOL_POLICY_CONFIDENTIALITY_KEY, ConfidentialityLabel) + max_allowed_confidentiality = parse_enum( + _TOOL_POLICY_MAX_ALLOWED_CONFIDENTIALITY_KEY, + ConfidentialityLabel, + ) + + accepts_untrusted = False + if _TOOL_POLICY_ACCEPTS_UNTRUSTED_KEY in props: + accepts_untrusted_raw = props[_TOOL_POLICY_ACCEPTS_UNTRUSTED_KEY] + if type(accepts_untrusted_raw) is bool: + accepts_untrusted = accepts_untrusted_raw + else: + errors.append(f"{_TOOL_POLICY_ACCEPTS_UNTRUSTED_KEY} must be a boolean") + + principal_state = "absent" + principals: tuple[_Principal, ...] | None = None + if PRINCIPAL_METADATA_KEY in props: + try: + principals = _canonical_principals( + props[PRINCIPAL_METADATA_KEY], + source=f"tool {function.name}", + ) + except ValueError: + principal_state = "invalid" + else: + principal_state = "valid" + + return _ToolSecurityPolicy( + source_integrity=cast(IntegrityLabel | None, source_integrity), + confidentiality=cast(ConfidentialityLabel | None, confidentiality) or default_confidentiality, + max_allowed_confidentiality=cast(ConfidentialityLabel | None, max_allowed_confidentiality), + accepts_untrusted=accepts_untrusted, + principal_state=principal_state, + principals=principals, + internal_security_tool=props.get(_INTERNAL_SECURITY_TOOL) is _INTERNAL_SECURITY_TOOL_MARKER, + errors=tuple(errors), + ) + + def check_confidentiality_allowed( context_label: ContentLabel, max_allowed: ConfidentialityLabel, @@ -885,8 +1020,32 @@ def _strict_json_value( if not math.isfinite(value): raise ValueError(f"{path} contains a non-finite float") return value + if canonical and type(value) is bytes: + return ["bytes", value.hex()] if canonical and isinstance(value, AgentSession): return ["agent_session", value.session_id] + if canonical and isinstance(value, VariableReferenceContent): + return [ + "variable_reference", + _strict_json_value( + value.to_dict(), + path=path, + canonical=True, + active_container_ids=active_container_ids, + ), + ] + if canonical and isinstance(value, BaseModel): + value_type = type(value) + return [ + "pydantic", + f"{value_type.__module__}.{value_type.__qualname__}", + _strict_json_value( + value.model_dump(exclude_unset=True), + path=path, + canonical=True, + active_container_ids=active_container_ids, + ), + ] if allow_content and isinstance(value, Content): return _strict_json_value( value.to_dict(), @@ -956,6 +1115,57 @@ def _strict_json_value( raise TypeError(f"{path} contains unsupported type {type(value).__name__}") +def _security_invocation_authority_token( + function: FunctionTool, + *, + tool_policy: _ToolSecurityPolicy, + session: AgentSession | None, + security_scope: _SecurityScope, + source_arguments: Any, + source_kwargs: Mapping[str, Any], + resolved_arguments: Any, + resolved_kwargs: Mapping[str, Any], + labels: Sequence[ContentLabel], +) -> _SecurityInvocationAuthority: + """Bind inspected references, resolved values, and policy labels without validators.""" + arguments_token = _argument_authority_token( + { + "name": function.name, + "source_arguments": source_arguments, + "source_kwargs": dict(source_kwargs), + "resolved_arguments": resolved_arguments, + "resolved_kwargs": dict(resolved_kwargs), + "labels": [label.to_dict() for label in labels], + }, + boundary="security policy", + ) + return _SecurityInvocationAuthority( + function=function, + tool_policy=tool_policy, + session=session, + security_scope=security_scope, + arguments_token=arguments_token, + ) + + +def _bindable_runtime_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]: + """Normalize runtime kwargs using the same canonical identities as approval binding.""" + bindable: dict[str, Any] = {} + for key, value in kwargs.items(): + try: + bindable[key] = _strict_json_value( + value, + path=f"security policy runtime kwarg {key!r}", + canonical=True, + ) + except (RecursionError, TypeError, ValueError) as exc: + raise MiddlewareFailure( + "Cannot safely bind security policy to opaque runtime authority. " + "Use JSON-native values, bytes, or a supported framework identity." + ) from exc + return bindable + + def _durable_state_snapshot( value: Any, *, @@ -1022,7 +1232,7 @@ def _durable_state_snapshot( def _encode_durable_state(value: Any, *, path: str, allow_content: bool = False) -> str: """Encode user-controlled state so AgentSession cannot reinterpret type discriminators.""" snapshot = _durable_state_snapshot(value, path=path, allow_content=allow_content) - return json.dumps(snapshot, ensure_ascii=False, allow_nan=False, separators=(",", ":")) + return json.dumps(snapshot, ensure_ascii=True, allow_nan=False, separators=(",", ":")) def _decode_durable_state(value: Any, *, path: str) -> Any: @@ -1460,12 +1670,12 @@ def _degrade_rewritten_arguments(context: FunctionInvocationContext) -> None: rewritten = context.metadata.get(_REWRITTEN_ARGUMENT_INDICES_KEY) if not rewritten: return - final_keys = _top_level_argument_keys(context) + final_keys = _top_level_argument_keys(context.arguments, context.kwargs) pre_validation_keys = set(cast(dict[str, set[int]], rewritten).keys()) iterable_keys = final_keys if pre_validation_keys != final_keys else pre_validation_keys degraded: dict[str, set[int]] = {} for arg_name in iterable_keys: - value, _ = _top_level_argument_value(context, arg_name) + value, _ = _top_level_argument_value(context.arguments, context.kwargs, arg_name) if isinstance(value, (list, tuple)): degraded[arg_name] = set(range(len(cast(Sequence[Any], value)))) else: @@ -1667,30 +1877,37 @@ def _resolve_value( ) return value - def _expand_variable_references_in_context(self, context: FunctionInvocationContext) -> list[ContentLabel]: - """Expand owned references in invocation values and return their stored labels.""" - if context.function.name in _VARIABLE_ID_CONSUMERS: - return [] - + def _resolve_invocation_arguments( + self, + context: FunctionInvocationContext, + *, + tool_policy: _ToolSecurityPolicy | None = None, + ) -> _ResolvedInvocationArguments: + """Resolve hidden references into detached values without invoking validators.""" labels: list[ContentLabel] = [] active_variables: set[str] = set() reference_count = [0] rewritten_paths: set[tuple[str | int, ...]] = set() kwargs_rewritten_paths: set[tuple[str | int, ...]] = set() - if context.arguments: - context.arguments = self._resolve_value( - context.arguments, + + source_arguments = context.arguments + source_kwargs = dict(context.kwargs) + arguments = source_arguments + kwargs = dict(source_kwargs) + if context.function.name not in _VARIABLE_ID_CONSUMERS and arguments: + arguments = self._resolve_value( + arguments, labels, depth=0, active_variables=active_variables, reference_count=reference_count, rewritten_paths=rewritten_paths, ) - if context.kwargs: - context.kwargs = cast( + if context.function.name not in _VARIABLE_ID_CONSUMERS and kwargs: + kwargs = cast( dict[str, Any], self._resolve_value( - context.kwargs, + kwargs, labels, depth=0, active_variables=active_variables, @@ -1699,7 +1916,7 @@ def _expand_variable_references_in_context(self, context: FunctionInvocationCont ), ) - args_mapping = cast(Any, context.arguments) + args_mapping = cast(Any, arguments) arg_names_in_arguments: set[str] = ( set(cast(dict[str, Any], args_mapping).keys()) if isinstance(args_mapping, Mapping) else set() ) @@ -1716,7 +1933,7 @@ def _expand_variable_references_in_context(self, context: FunctionInvocationCont if arg_name not in rewritten_args: rewritten_args[arg_name] = set() - arg_value, arg_source = _top_level_argument_value(context, arg_name) + arg_value, arg_source = _top_level_argument_value(arguments, kwargs, arg_name) if arg_source is None: rewritten_args[arg_name].add(-1) @@ -1736,8 +1953,112 @@ def _expand_variable_references_in_context(self, context: FunctionInvocationCont else: rewritten_args[arg_name].add(-1) - context.metadata[_REWRITTEN_ARGUMENT_INDICES_KEY] = rewritten_args - return labels + source_security_kwargs = _bindable_runtime_kwargs(source_kwargs) + security_kwargs = _bindable_runtime_kwargs(kwargs) + resolved_tool_policy = tool_policy or _parse_tool_security_policy( + context.function, + default_confidentiality=self.default_confidentiality, + ) + authority_token = _security_invocation_authority_token( + context.function, + tool_policy=resolved_tool_policy, + session=context.session, + security_scope=self._scope, + source_arguments=source_arguments, + source_kwargs=source_security_kwargs, + resolved_arguments=arguments, + resolved_kwargs=security_kwargs, + labels=labels, + ) + return _ResolvedInvocationArguments( + arguments=arguments, + kwargs=kwargs, + labels=tuple(labels), + rewritten_arguments=rewritten_args, + authority_token=authority_token, + ) + + def _ensure_preflight_authority_unchanged( + self, + context: FunctionInvocationContext, + ) -> _ResolvedInvocationArguments: + """Reject argument or runtime-kwarg changes after detached security preflight.""" + resolved = self._resolve_invocation_arguments(context) + expected_authority = context.metadata.get(_SECURITY_RESOLVED_AUTHORITY_CONTEXT_KEY) + if expected_authority is None: + return resolved + if resolved.authority_token != expected_authority: + raise MiddlewareFailure( + "Function arguments changed after security middleware processed them. " + "Install argument-repair middleware before security middleware." + ) + return resolved + + def _materialized_authority_token(self, context: FunctionInvocationContext) -> _SecurityInvocationAuthority: + """Bind one fully materialized execution attempt.""" + tool_policy = _parse_tool_security_policy( + context.function, + default_confidentiality=self.default_confidentiality, + ) + arguments_token = _argument_authority_token( + { + "name": context.function.name, + "arguments": context.arguments, + "kwargs": _bindable_runtime_kwargs(context.kwargs), + }, + boundary="materialized security policy", + ) + return _SecurityInvocationAuthority( + function=context.function, + tool_policy=tool_policy, + session=context.session, + security_scope=self._scope, + arguments_token=arguments_token, + ) + + def _materialize_arguments(self, context: FunctionInvocationContext) -> None: + """Apply preflighted hidden values and validate only after policy allows the call.""" + if context.metadata.get(_SECURITY_ARGUMENTS_MATERIALIZED_CONTEXT_KEY) is True: + expected_materialized_authority = context.metadata.get(_SECURITY_MATERIALIZED_AUTHORITY_CONTEXT_KEY) + if ( + expected_materialized_authority is not None + and self._materialized_authority_token(context) != expected_materialized_authority + ): + raise MiddlewareFailure( + "Function arguments changed after security middleware processed them. " + "Install argument-repair middleware before security middleware." + ) + return + + policy_errors = context.metadata.get(_SECURITY_POLICY_ERRORS_CONTEXT_KEY) + if policy_errors: + raise MiddlewareFailure( + f"Invalid FIDES policy for tool '{context.function.name}': {'; '.join(policy_errors)}" + ) + + resolved = self._ensure_preflight_authority_unchanged(context) + + context.arguments = resolved.arguments + context.kwargs = resolved.kwargs + context.metadata[_REWRITTEN_ARGUMENT_INDICES_KEY] = resolved.rewritten_arguments + context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = _argument_authority_token( + context.arguments, + boundary="security policy", + ) + if context.metadata.get(_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY) is True: + pre_validation_token = _argument_authority_token(context.arguments, boundary="security policy") + context.function._prepare_context_arguments( # pyright: ignore[reportPrivateUsage] + context, + context.arguments, + ) + post_validation_token = _argument_authority_token(context.arguments, boundary="security policy") + context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = post_validation_token + if pre_validation_token != post_validation_token: + if context.metadata.get("user_approved_violation") is True: + raise MiddlewareFailure("Function arguments changed during validation after FIDES policy approval.") + self._degrade_rewritten_arguments(context) + context.metadata[_SECURITY_MATERIALIZED_AUTHORITY_CONTEXT_KEY] = self._materialized_authority_token(context) + context.metadata[_SECURITY_ARGUMENTS_MATERIALIZED_CONTEXT_KEY] = True def _get_input_labels(self, context: FunctionInvocationContext) -> list[ContentLabel]: """Extract security labels from tool input arguments. @@ -1817,18 +2138,10 @@ def _get_source_integrity(self, context: FunctionInvocationContext) -> Integrity Returns: IntegrityLabel if declared, None if not declared. """ - function_props = _get_additional_properties(context.function) - source_integrity_str = function_props.get("source_integrity", None) - - if source_integrity_str is not None: - try: - return IntegrityLabel(source_integrity_str) - except ValueError: - logger.warning( - f"Invalid source_integrity '{source_integrity_str}' for function " - f"'{context.function.name}', ignoring" - ) - return None + return _parse_tool_security_policy( + context.function, + default_confidentiality=self.default_confidentiality, + ).source_integrity # ========== Helper utilities ========== @@ -1893,10 +2206,11 @@ async def process( context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]], ) -> None: - """Resolve hidden arguments, publish their labels, and label the result.""" + """Preflight hidden arguments, publish their labels, and label the result.""" scope_token = self._activate_security_scope(context) middleware_token = _current_middleware.set(self) context_token = _current_context.set(context) + materialization_token: Token[tuple[Any, Callable[[Any], None], asyncio.Task[Any] | None] | None] | None = None try: function_name = context.function.name if "original_arguments_for_messages" not in context.metadata: @@ -1907,35 +2221,37 @@ async def process( context.metadata["security_original_runtime_kwargs"] = dict(context.kwargs) else: context.kwargs = dict(cast(dict[str, Any], context.metadata["security_original_runtime_kwargs"])) - + context.metadata.pop(_SECURITY_ARGUMENTS_MATERIALIZED_CONTEXT_KEY, None) + context.metadata.pop(_SECURITY_ARGUMENTS_REQUIRE_MATERIALIZATION_CONTEXT_KEY, None) + context.metadata.pop(_SECURITY_MATERIALIZED_AUTHORITY_CONTEXT_KEY, None) + context.metadata.pop(_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY, None) + context.metadata.pop(_SECURITY_RESOLVED_AUTHORITY_CONTEXT_KEY, None) + + tool_policy = _parse_tool_security_policy( + context.function, + default_confidentiality=self.default_confidentiality, + ) + context.metadata[_SECURITY_POLICY_ERRORS_CONTEXT_KEY] = tool_policy.errors input_labels = self._get_input_labels(context) - declared_source_integrity = self._get_source_integrity(context) - confidentiality = self._get_function_confidentiality(context) + declared_source_integrity = tool_policy.source_integrity + confidentiality = tool_policy.confidentiality standing_guidance_snapshot = self._get_standing_guidance(context.function) - # Expand hidden references before execution and retain their stored labels. - resolved_labels = self._expand_variable_references_in_context(context) - context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = _argument_authority_token( - context.arguments, - boundary="security policy", + # Resolve into detached values for label and approval preflight only. + resolved = self._resolve_invocation_arguments(context, tool_policy=tool_policy) + resolved_labels = list(resolved.labels) + context.metadata[_REWRITTEN_ARGUMENT_INDICES_KEY] = resolved.rewritten_arguments + context.metadata[_SECURITY_ARGUMENTS_REQUIRE_MATERIALIZATION_CONTEXT_KEY] = bool( + resolved.rewritten_arguments ) - if context.metadata.get(_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY) is True: - pre_validation_token = _argument_authority_token(context.arguments, boundary="security policy") - context.function._prepare_context_arguments( # pyright: ignore[reportPrivateUsage] - context, - context.arguments, - ) - post_validation_token = _argument_authority_token(context.arguments, boundary="security policy") - context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = post_validation_token - if pre_validation_token != post_validation_token: - self._degrade_rewritten_arguments(context) + context.metadata[_SECURITY_RESOLVED_AUTHORITY_CONTEXT_KEY] = resolved.authority_token argument_labels = [*input_labels, *resolved_labels] argument_label = combine_labels(*argument_labels) if argument_labels else ContentLabel() declared_confidentiality_label = ContentLabel( confidentiality=confidentiality, - metadata=self._get_function_principal_metadata(context, confidentiality), + metadata=self._get_function_principal_metadata(context.function, tool_policy), ) result_policy_label = combine_labels(declared_confidentiality_label, argument_label) result_confidentiality = result_policy_label.confidentiality @@ -1978,11 +2294,25 @@ async def process( context.metadata["argument_label"] = argument_label context.metadata["effective_invocation_label"] = combine_labels(context_label, argument_label) - await call_next() + materialization_token = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.set(( + context, + self._materialize_arguments, + asyncio.current_task(), + )) + try: + await call_next() + except MiddlewareTermination: + if context.metadata.get(_SECURITY_ARGUMENTS_MATERIALIZED_CONTEXT_KEY) is not True: + self._ensure_preflight_authority_unchanged(context) + raise + if context.metadata.get(_SECURITY_ARGUMENTS_MATERIALIZED_CONTEXT_KEY) is not True: + self._ensure_preflight_authority_unchanged(context) if isinstance(context.result, Content) and context.result.type == "function_approval_request": return self._label_result(context, function_name, fallback_label, standing_guidance_snapshot) finally: + if materialization_token is not None: + _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.reset(materialization_token) _current_context.reset(context_token) _current_middleware.reset(middleware_token) self._active_security_scope.reset(scope_token) @@ -2085,39 +2415,23 @@ def _get_function_confidentiality(self, context: FunctionInvocationContext) -> C Returns: The confidentiality label for this function. """ - # Check function's additional_properties for confidentiality setting - function_props = _get_additional_properties(context.function) - confidentiality_str = function_props.get("confidentiality", None) - - if confidentiality_str: - try: - return ConfidentialityLabel(confidentiality_str) - except ValueError: - logger.warning( - f"Invalid confidentiality label '{confidentiality_str}' " - f"for function '{context.function.name}', using default" - ) - - return self.default_confidentiality + return _parse_tool_security_policy( + context.function, + default_confidentiality=self.default_confidentiality, + ).confidentiality def _get_function_principal_metadata( self, - context: FunctionInvocationContext, - confidentiality: ConfidentialityLabel, + function: FunctionTool, + tool_policy: _ToolSecurityPolicy, ) -> dict[str, Any]: """Read a locally declared source principal set.""" - if confidentiality != ConfidentialityLabel.USER_IDENTITY: + if tool_policy.confidentiality != ConfidentialityLabel.USER_IDENTITY: return {} - function_props = _get_additional_properties(context.function) - try: - principals = _canonical_principals( - function_props.get(PRINCIPAL_METADATA_KEY), - source=f"tool {context.function.name}", - ) - except ValueError as exc: - logger.warning("Invalid USER_IDENTITY source principals for tool '%s': %s", context.function.name, exc) + if tool_policy.principal_state != "valid" or tool_policy.principals is None: + logger.warning("Invalid USER_IDENTITY source principals for tool '%s'.", function.name) return {} - return {PRINCIPAL_METADATA_KEY: _principal_list(principals)} + return {PRINCIPAL_METADATA_KEY: _principal_list(tool_policy.principals)} def _process_result_with_embedded_labels( self, @@ -2533,11 +2847,11 @@ def rewritten_arguments(context: FunctionInvocationContext | None = None) -> dic current_snapshot = _argument_authority_token(context.arguments, boundary="security policy") if current_snapshot != indices_snapshot: rewritten_keys = set(cast(dict[str, set[int]], rewritten).keys()) - final_keys = _top_level_argument_keys(context) + final_keys = _top_level_argument_keys(context.arguments, context.kwargs) iterable_keys = final_keys if rewritten_keys != final_keys else rewritten_keys result: dict[str, set[int]] = {} for arg_name in iterable_keys: - value, _ = _top_level_argument_value(context, arg_name) + value, _ = _top_level_argument_value(context.arguments, context.kwargs, arg_name) if isinstance(value, (list, tuple)): result[arg_name] = set(range(len(cast(Sequence[Any], value)))) else: @@ -2826,6 +3140,13 @@ def _resolved_arguments(self, context: FunctionInvocationContext) -> dict[str, A return dict(context.arguments) def _resolved_call_signature(self, context: FunctionInvocationContext) -> str: + tracker = _current_middleware.get() + if tracker is not None: + resolved = tracker._resolve_invocation_arguments(context) # pyright: ignore[reportPrivateUsage] + return self._signature_from_parts( + context.function.name, + {"arguments": resolved.arguments, "kwargs": resolved.kwargs}, + ) return self._signature_from_parts( context.function.name, {"arguments": self._resolved_arguments(context), "kwargs": dict(context.kwargs)}, @@ -2860,10 +3181,12 @@ def _label_key(label_data: Any) -> str: return "/" def _destination_principal_key(self, context: FunctionInvocationContext) -> str: - function_props = _get_additional_properties(context.function) - if function_props.get("max_allowed_confidentiality") != ConfidentialityLabel.USER_IDENTITY.value: + tool_policy = _parse_tool_security_policy(context.function) + if tool_policy.max_allowed_confidentiality != ConfidentialityLabel.USER_IDENTITY: return "" - return _principal_binding_key(function_props.get(PRINCIPAL_METADATA_KEY), source="approval destination") + if tool_policy.principal_state != "valid" or tool_policy.principals is None: + raise ValueError("USER_IDENTITY destination principals are missing or invalid") + return _principal_binding_key(_principal_list(tool_policy.principals), source="approval destination") def _session_key(self, context: FunctionInvocationContext) -> str: return context.session.session_id if context.session is not None else "" @@ -3064,6 +3387,44 @@ def _block_unsafe_approval_binding( } raise MiddlewareTermination("Unsafe policy approval binding") + def _block_invalid_tool_policy( + self, + context: FunctionInvocationContext, + *, + context_label: ContentLabel, + errors: Sequence[str], + ) -> NoReturn: + reason = "; ".join(errors) + self._log_violation({ + "type": "invalid_policy", + "function": context.function.name, + "context_label": context_label.to_dict(), + "turn": context.metadata.get("turn_number", -1), + "reason": reason, + }) + context.result = { + "error": f"Invalid FIDES tool policy: {reason}", + "function": context.function.name, + "context_label": context_label.to_dict(), + "violation_type": "invalid_policy", + } + raise MiddlewareTermination("Invalid FIDES tool policy") + + @staticmethod + def _ensure_terminal_materialization_boundary(context: FunctionInvocationContext) -> None: + """Reject allowed hidden arguments when no terminal pipeline boundary can materialize them.""" + if context.metadata.get(_SECURITY_ARGUMENTS_REQUIRE_MATERIALIZATION_CONTEXT_KEY) is not True: + return + active_pipeline = _ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT.get() + if ( + active_pipeline is None + or active_pipeline[0] is not context + or active_pipeline[1] is not asyncio.current_task() + ): + raise MiddlewareFailure( + "Hidden function arguments require FunctionMiddlewarePipeline for terminal materialization." + ) + async def process( self, context: FunctionInvocationContext, @@ -3102,10 +3463,14 @@ async def _process_in_scope( argument_label = self._resolve_label(context.metadata.get("argument_label")) effective_label = combine_labels(context_label, argument_label) context.metadata["effective_invocation_label"] = effective_label - function_props = _get_additional_properties(context.function) - accepts_untrusted = ( - function_name in self.allow_untrusted_tools or function_props.get("accepts_untrusted") is True - ) + tool_policy = _parse_tool_security_policy(context.function) + if tool_policy.errors: + self._block_invalid_tool_policy( + context, + context_label=context_label, + errors=tool_policy.errors, + ) + accepts_untrusted = function_name in self.allow_untrusted_tools or tool_policy.accepts_untrusted violations: list[dict[str, Any]] = [] if context_label.integrity == IntegrityLabel.UNTRUSTED and not accepts_untrusted: @@ -3169,6 +3534,7 @@ async def _process_in_scope( }) if not violations: + self._ensure_terminal_materialization_boundary(context) await call_next() return for violation in violations: @@ -3206,6 +3572,7 @@ async def _process_in_scope( else: logger.warning("WARNING: Tool '%s' policy violation(s) [%s] (allowed)", function_name, disclosed) + self._ensure_terminal_materialization_boundary(context) await call_next() @staticmethod @@ -3256,7 +3623,13 @@ def _check_confidentiality_policy_detailed( Returns: Dict with keys: passed (bool), failure_type (str), reason (str). """ - function_props = _get_additional_properties(context.function) + tool_policy = _parse_tool_security_policy(context.function) + if tool_policy.errors: + return { + "passed": False, + "failure_type": "invalid_policy", + "reason": "; ".join(tool_policy.errors), + } conf_hierarchy = { ConfidentialityLabel.PUBLIC: 0, @@ -3267,27 +3640,16 @@ def _check_confidentiality_policy_detailed( # Check max_allowed_confidentiality (output restriction / data exfiltration prevention) # Context confidentiality must be <= max allowed level # This prevents PRIVATE data from being written to PUBLIC destinations - max_allowed_conf = function_props.get("max_allowed_confidentiality", None) - if max_allowed_conf is not None: - try: - max_allowed_level = ConfidentialityLabel(max_allowed_conf) - if conf_hierarchy[label.confidentiality] > conf_hierarchy[max_allowed_level]: - return { - "passed": False, - "failure_type": "max_allowed_confidentiality", - "reason": ( - f"Cannot write {label.confidentiality.value.upper()} data to " - f"{max_allowed_level.value.upper()} destination (data exfiltration blocked)" - ), - } - except ValueError: - logger.warning(f"Invalid max_allowed_confidentiality: {max_allowed_conf}") - if label.confidentiality == ConfidentialityLabel.USER_IDENTITY: - return { - "passed": False, - "failure_type": "principal_mismatch", - "reason": "USER_IDENTITY destination policy is invalid", - } + max_allowed_level = tool_policy.max_allowed_confidentiality + if max_allowed_level is not None and conf_hierarchy[label.confidentiality] > conf_hierarchy[max_allowed_level]: + return { + "passed": False, + "failure_type": "max_allowed_confidentiality", + "reason": ( + f"Cannot write {label.confidentiality.value.upper()} data to " + f"{max_allowed_level.value.upper()} destination (data exfiltration blocked)" + ), + } if label.confidentiality == ConfidentialityLabel.USER_IDENTITY: try: @@ -3297,21 +3659,20 @@ def _check_confidentiality_policy_detailed( source="source label", ) ) - if function_props.get(_INTERNAL_SECURITY_TOOL) is _INTERNAL_SECURITY_TOOL_MARKER: + if tool_policy.internal_security_tool: destination_principals = source_principals else: - if max_allowed_conf != ConfidentialityLabel.USER_IDENTITY.value: + if ( + max_allowed_level != ConfidentialityLabel.USER_IDENTITY + or tool_policy.principal_state != "valid" + or tool_policy.principals is None + ): return { "passed": False, "failure_type": "principal_mismatch", "reason": "USER_IDENTITY destination does not declare an authorized principal set", } - destination_principals = set( - _canonical_principals( - function_props.get(PRINCIPAL_METADATA_KEY), - source=f"tool {context.function.name}", - ) - ) + destination_principals = set(tool_policy.principals) except ValueError: return { "passed": False, diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 32a128775e..a0163d6065 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -7,9 +7,11 @@ import logging import math import warnings +from collections.abc import Awaitable, Callable from datetime import datetime, timedelta, timezone +from pathlib import Path from types import MappingProxyType, SimpleNamespace -from typing import Annotated, Any, cast +from typing import Annotated, Any, Literal, cast from unittest.mock import AsyncMock import pytest @@ -20,6 +22,7 @@ AgentSession, ChatResponse, ExperimentalFeature, + FileSessionStore, FunctionInvocationContext, FunctionMiddleware, Message, @@ -825,7 +828,7 @@ async def next_fn(): @pytest.mark.asyncio async def test_bracketed_variable_reference_expanded_before_call_next(self, middleware): - """Bracketed variable placeholders should be expanded before tool execution.""" + """Bracketed variable placeholders should be expanded at terminal execution.""" class MessageArgs(BaseModel): summary: str @@ -850,16 +853,16 @@ async def send_message(summary: str) -> str: arguments=MessageArgs(summary=f"[{variable_id}]"), ) - async def next_fn() -> None: + async def next_fn() -> list[Content]: current_args = context.arguments if isinstance(current_args, BaseModel): summary = current_args.model_dump()["summary"] else: summary = current_args["summary"] assert summary == expected_summary - context.result = [Content.from_text("sent")] + return [Content.from_text("sent")] - await middleware.process(context, next_fn) + await FunctionMiddlewarePipeline(middleware).execute(context, lambda _: next_fn()) async def test_transitive_variable_reference_expands_and_collects_all_labels(self, middleware) -> None: store = middleware.get_variable_store() @@ -874,12 +877,12 @@ async def test_transitive_variable_reference_expands_and_collects_all_labels(sel function = self._variable_sink() context = FunctionInvocationContext(function=function, arguments={"payload": f"[{root_id}]"}) - async def next_fn() -> None: + async def next_fn() -> list[Content]: assert context.arguments == {"payload": "expanded value"} assert context.metadata["argument_label"].integrity == IntegrityLabel.UNTRUSTED - context.result = [Content.from_text("done")] + return [Content.from_text("done")] - await middleware.process(context, next_fn) + await FunctionMiddlewarePipeline(middleware).execute(context, lambda _: next_fn()) async def test_three_level_variable_reference_chain_expands(self, middleware) -> None: store = middleware.get_variable_store() @@ -889,11 +892,11 @@ async def test_three_level_variable_reference_chain_expands(self, middleware) -> function = self._variable_sink() context = FunctionInvocationContext(function=function, arguments={"payload": f"[{first_id}]"}) - async def next_fn() -> None: + async def next_fn() -> list[Content]: assert context.arguments == {"payload": "third"} - context.result = [Content.from_text("done")] + return [Content.from_text("done")] - await middleware.process(context, next_fn) + await FunctionMiddlewarePipeline(middleware).execute(context, lambda _: next_fn()) @pytest.mark.parametrize("indirect", [False, True], ids=["direct", "indirect"]) async def test_variable_reference_cycle_fails_before_tool_execution(self, middleware, indirect: bool) -> None: @@ -932,7 +935,7 @@ async def test_nested_mixed_trust_references_expand_and_taint_arguments(self, mi }, ) - async def next_fn() -> None: + async def next_fn() -> list[Content]: assert context.arguments == { "payload": { "trusted": "trusted", @@ -940,9 +943,9 @@ async def next_fn() -> None: } } assert context.metadata["argument_label"].integrity == IntegrityLabel.UNTRUSTED - context.result = [Content.from_text("done")] + return [Content.from_text("done")] - await middleware.process(context, next_fn) + await FunctionMiddlewarePipeline(middleware).execute(context, lambda _: next_fn()) async def test_variable_reference_depth_limit_fails_before_tool_execution(self, middleware) -> None: store = middleware.get_variable_store() @@ -976,7 +979,7 @@ async def next_fn() -> None: @pytest.mark.asyncio async def test_json_string_variable_reference_expands_only_response_before_call_next(self, middleware): - """JSON-serialized hidden payloads should expose only the response text to tools.""" + """JSON-serialized hidden payloads should expose only the response text at execution.""" class MessageArgs(BaseModel): summary: str @@ -1012,7 +1015,7 @@ async def send_message(summary: str) -> str: arguments=MessageArgs(summary=f"Security review complete. [{variable_id}]"), ) - async def next_fn() -> None: + async def next_fn() -> list[Content]: current_args = context.arguments if isinstance(current_args, BaseModel): summary = current_args.model_dump()["summary"] @@ -1020,9 +1023,9 @@ async def next_fn() -> None: summary = current_args["summary"] assert summary == f"Security review complete. {response_text}" assert '"response"' not in summary - context.result = [Content.from_text("sent")] + return [Content.from_text("sent")] - await middleware.process(context, next_fn) + await FunctionMiddlewarePipeline(middleware).execute(context, lambda _: next_fn()) async def test_json_response_object_is_not_mistaken_for_quarantine_payload(self, middleware) -> None: """Ordinary JSON objects containing ``response`` survive expansion intact.""" @@ -1057,13 +1060,13 @@ async def send_message(summary: str) -> str: arguments=MessageArgs(summary=f"[{variable_id}]"), ) - async def next_fn() -> None: + async def next_fn() -> list[Content]: current_args = context.arguments assert isinstance(current_args, dict) assert current_args["summary"] == stored_payload - context.result = [Content.from_text("sent")] + return [Content.from_text("sent")] - await middleware.process(context, next_fn) + await FunctionMiddlewarePipeline(middleware).execute(context, lambda _: next_fn()) class TestPolicyEnforcementMiddleware: @@ -1105,6 +1108,65 @@ async def next_fn(): assert context.result == [Content.from_text("mock result")] + @pytest.mark.parametrize( + "properties", + [ + {"source_integrity": "TRUSTED"}, + {"source_integrity": None}, + {"confidentiality": "PUBLIC"}, + {"confidentiality": None}, + {"max_allowed_confidentiality": "PUBLIC"}, + {"max_allowed_confidentiality": None}, + {"accepts_untrusted": "true"}, + {"accepts_untrusted": None}, + {"max_allowed_confidentialty": "public"}, + {"Max_Allowed_Confidentiality": "public"}, + ], + ) + async def test_invalid_tool_policy_is_blocked(self, mock_function, properties: dict[str, Any]) -> None: + """Unknown or mistyped FIDES metadata never removes a security restriction.""" + mock_function.additional_properties = properties + middleware = PolicyEnforcementFunctionMiddleware(block_on_violation=False) + context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + ) + context.metadata["context_label"] = ContentLabel() + executed = False + + async def next_fn() -> None: + nonlocal executed + executed = True + + with pytest.raises(MiddlewareTermination): + await middleware.process(context, next_fn) + + assert executed is False + assert context.result["violation_type"] == "invalid_policy" + assert middleware.get_audit_log()[-1]["type"] == "invalid_policy" + + async def test_unrelated_tool_metadata_is_allowed(self, mock_function) -> None: + """Unrecognized non-FIDES metadata remains available to tool integrations.""" + mock_function.additional_properties = { + "confidentiality_note": "integration-owned metadata", + "provider_feature": True, + } + middleware = PolicyEnforcementFunctionMiddleware() + context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + ) + context.metadata["context_label"] = ContentLabel() + executed = False + + async def next_fn() -> None: + nonlocal executed + executed = True + + await middleware.process(context, next_fn) + + assert executed is True + @pytest.mark.asyncio async def test_untrusted_call_blocked(self, middleware, mock_function): """Test that untrusted tool calls are blocked.""" @@ -3525,6 +3587,33 @@ async def test_session_state_is_detached_json_and_rejects_unsupported_values(sel with pytest.raises(ValueError, match="integrity and confidentiality"): malformed_tracker.get_context_label() + @pytest.mark.parametrize("serialization_format", ["json", "msgpack"]) + async def test_hidden_surrogate_round_trips_in_file_session_store( + self, + tmp_path: Path, + serialization_format: Literal["json", "msgpack"], + ) -> None: + """Opaque FIDES state remains serializable when hidden text contains a lone surrogate.""" + config = SecureAgentConfig() + session = AgentSession(session_id=f"surrogate-{serialization_format}") + tracker, _ = await _get_session_security_middleware(config, session) + variable_id = tracker.get_variable_store().store( + "\ud800", + ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ), + ) + store = FileSessionStore(tmp_path, serialization_format=serialization_format) + + await store.set(session.session_id, session) + restored = await store.get(session.session_id) + + assert restored is not None + restored_content, restored_label = config.get_variable_store(restored).retrieve(variable_id) + assert restored_content == "\ud800" + assert restored_label.confidentiality == ConfidentialityLabel.PRIVATE + async def test_audit_and_pending_approval_are_session_scoped_and_resume_after_restore(self) -> None: """Audit and approval state persists only in the owning session.""" config = SecureAgentConfig(approval_on_violation=True) @@ -5439,12 +5528,12 @@ async def transform(data: dict[str, Any]) -> str: data["value"] = f"[{variable_id}]" context = FunctionInvocationContext(function=function, arguments=Args(data=data)) - async def next_fn() -> None: + async def next_fn() -> list[Content]: assert isinstance(context.arguments, dict) assert context.arguments["data"]["value"] == "input" - context.result = [Content.from_text("transformed")] + return [Content.from_text("transformed")] - await tracker.process(context, next_fn) + await FunctionMiddlewarePipeline(tracker).execute(context, lambda _: next_fn()) expected_integrity = stored_integrity if stored_integrity is not None else default_integrity if claimed_integrity == IntegrityLabel.UNTRUSTED: @@ -5533,12 +5622,12 @@ async def transform(value: str) -> str: arguments=Args(value=f"[{variable_id}]"), ) - async def next_fn() -> None: + async def next_fn() -> list[Content]: current_args = context.arguments assert isinstance(current_args, dict) - context.result = [Content.from_text(current_args["value"].upper())] + return [Content.from_text(current_args["value"].upper())] - await middleware.process(context, next_fn) + await FunctionMiddlewarePipeline(middleware).execute(context, lambda _: next_fn()) assert context.metadata["result_label"].integrity == IntegrityLabel.TRUSTED assert context.metadata["result_label"].confidentiality == confidentiality @@ -7543,7 +7632,7 @@ async def sink(value: Any) -> str: ) @pytest.mark.parametrize( - ("make_value", "expected"), + ("make_value", "_expected"), [ pytest.param(lambda variable_id: f"[{variable_id}]", "payload", id="bracketed"), pytest.param(lambda variable_id: f" [ {variable_id} ] ", "payload", id="whitespace"), @@ -7562,7 +7651,7 @@ async def sink(value: Any) -> str: async def test_all_reference_forms_resolve_and_block_untrusted_arguments( self, make_value: Any, - expected: Any, + _expected: Any, ) -> None: tracker = LabelTrackingFunctionMiddleware() policy = PolicyEnforcementFunctionMiddleware() @@ -7582,150 +7671,974 @@ async def test_all_reference_forms_resolve_and_block_untrusted_arguments( lambda _context: pytest.fail("Blocked tool must not execute"), ) - assert cast(dict[str, Any], context.arguments)["value"] == expected + assert cast(dict[str, Any], context.arguments)["value"] == value assert cast(dict[str, Any], context.metadata["original_arguments_for_messages"])["value"] == value assert context.metadata["argument_label"].integrity == IntegrityLabel.UNTRUSTED assert policy.get_audit_log()[-1]["type"] == "untrusted_arguments" - @pytest.mark.parametrize("value", ["[var_0123456789abcdef]", "var_0123456789abcdef"]) - async def test_unknown_references_remain_literal(self, value: str) -> None: + async def test_blocked_hidden_value_never_reaches_argument_validator(self) -> None: + """Policy preflight blocks before validators can observe expanded hidden content.""" tracker = LabelTrackingFunctionMiddleware() policy = PolicyEnforcementFunctionMiddleware() - received: list[str] = [] - context = FunctionInvocationContext( - function=self._sink(accepts_untrusted=False), - arguments={"value": value}, + secret = "validator-side-effect-secret" + variable_id = tracker.get_variable_store().store( + secret, + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), ) + observed: list[str] = [] + executed = False - async def execute(current: FunctionInvocationContext) -> list[Content]: - received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) - return [Content.from_text("sent")] + class SinkArgs(BaseModel): + value: str - await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + @field_validator("value") + @classmethod + def observe_value(cls, value: str) -> str: + observed.append(value) + return value - assert received == [value] - argument_label = cast(ContentLabel, context.metadata["argument_label"]) - assert argument_label.integrity == IntegrityLabel.TRUSTED - assert argument_label.confidentiality == ConfidentialityLabel.PUBLIC + async def sink(value: str) -> str: + nonlocal executed + executed = True + return value - async def test_accepts_untrusted_blind_forwards_without_model_context_taint(self) -> None: - tracker = LabelTrackingFunctionMiddleware() - policy = PolicyEnforcementFunctionMiddleware() - variable_id = tracker.get_variable_store().store( - "payload", - ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + sink_tool = FunctionTool( + func=sink, + name="validator_sink", + input_model=SinkArgs, + additional_properties={ + "accepts_untrusted": False, + "max_allowed_confidentiality": "private", + "source_integrity": "trusted", + }, ) - received: list[str] = [] - context = FunctionInvocationContext( - function=self._sink(accepts_untrusted=True), + function_call = Content.from_function_call( + call_id="validator-side-effect", + name=sink_tool.name, arguments={"value": f"[{variable_id}]"}, ) - async def execute(current: FunctionInvocationContext) -> list[Content]: - received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) - return [Content.from_text("sent")] - - await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + with pytest.raises(MiddlewareTermination): + await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) - assert received == ["payload"] - assert context.metadata["argument_label"].integrity == IntegrityLabel.UNTRUSTED - assert tracker.get_context_label().integrity == IntegrityLabel.TRUSTED + assert secret not in observed + assert executed is False - async def test_hidden_argument_resolution_does_not_require_reapproval(self) -> None: - """Security expansion preserves the approval-visible placeholder.""" + async def test_later_policy_blocks_before_hidden_value_validation(self) -> None: + """A permissive policy cannot materialize before a later policy rejects.""" tracker = LabelTrackingFunctionMiddleware() - policy = PolicyEnforcementFunctionMiddleware() + permissive_policy = PolicyEnforcementFunctionMiddleware(allow_untrusted_tools={"validator_sink"}) + strict_policy = PolicyEnforcementFunctionMiddleware() + secret = "later-policy-validator-secret" variable_id = tracker.get_variable_store().store( - "payload", + secret, ContentLabel(integrity=IntegrityLabel.UNTRUSTED), ) - sink = self._sink(accepts_untrusted=True) + observed: list[str] = [] + executed = False + + class SinkArgs(BaseModel): + value: str + + @field_validator("value") + @classmethod + def observe_value(cls, value: str) -> str: + observed.append(value) + return value + + async def sink(value: str) -> str: + nonlocal executed + executed = True + return value + + sink_tool = FunctionTool( + func=sink, + name="validator_sink", + input_model=SinkArgs, + additional_properties={ + "accepts_untrusted": False, + "max_allowed_confidentiality": "private", + "source_integrity": "trusted", + }, + ) function_call = Content.from_function_call( - call_id="approved-hidden-value", - id="approved-hidden-value-occurrence", - name=sink.name, + call_id="later-policy-validator-side-effect", + name=sink_tool.name, arguments={"value": f"[{variable_id}]"}, ) - approval_response = Content.from_function_approval_request( - id="approved-hidden-value-occurrence", - function_call=function_call, - ).to_function_approval_response(approved=True) - result = await _auto_invoke_function( - approval_response, - config=normalize_function_invocation_configuration(None), - tool_map={sink.name: sink}, - middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), - ) + with pytest.raises(MiddlewareTermination): + await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, permissive_policy, strict_policy), + ) - assert result.type == "function_result" - assert result.result == "payload" + assert secret not in observed + assert executed is False - async def test_hidden_argument_can_be_normalized_after_security_check(self) -> None: - """Final Pydantic coercion is not mistaken for post-policy mutation.""" + async def test_direct_composition_waits_for_all_policies_before_validation(self) -> None: + """Direct composition fails closed before any policy can expose hidden values.""" tracker = LabelTrackingFunctionMiddleware() - policy = PolicyEnforcementFunctionMiddleware() + permissive_policy = PolicyEnforcementFunctionMiddleware(allow_untrusted_tools={"validator_sink"}) + strict_policy = PolicyEnforcementFunctionMiddleware() + secret = "direct-composition-validator-secret" variable_id = tracker.get_variable_store().store( - "3", + secret, ContentLabel(integrity=IntegrityLabel.UNTRUSTED), ) + observed: list[str] = [] + executed = False - class StrictArgs(BaseModel): - value: int + class SinkArgs(BaseModel): + value: str - received: list[int] = [] + @field_validator("value") + @classmethod + def observe_value(cls, value: str) -> str: + observed.append(value) + return value - def strict_sink(value: int) -> str: - received.append(value) - return str(value) + async def sink(value: str) -> str: + nonlocal executed + executed = True + return value - strict_tool = FunctionTool( - func=strict_sink, - name="strict_sink", - input_model=StrictArgs, - additional_properties={"accepts_untrusted": True}, - ) - function_call = Content.from_function_call( - call_id="hidden-normalization", - name=strict_tool.name, - arguments={"value": f"[{variable_id}]"}, + sink_tool = FunctionTool( + func=sink, + name="validator_sink", + input_model=SinkArgs, + additional_properties={ + "accepts_untrusted": False, + "max_allowed_confidentiality": "private", + "source_integrity": "trusted", + }, ) + reference = f"[{variable_id}]" + context = FunctionInvocationContext(function=sink_tool, arguments={"value": reference}) + context.metadata[_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY] = True - result = await _auto_invoke_function( - function_call, - config=normalize_function_invocation_configuration(None), - tool_map={strict_tool.name: strict_tool}, - middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), - ) + async def execute() -> None: + nonlocal executed + executed = True - assert result.type == "function_result" - assert result.exception is None - assert received == [3] + with pytest.raises(MiddlewareFailure, match="require FunctionMiddlewarePipeline"): + await tracker.process( + context, + lambda: permissive_policy.process( + context, + lambda: strict_policy.process(context, execute), + ), + ) - @pytest.mark.parametrize("include_detailed_errors", [False, True]) - async def test_hidden_argument_validation_error_does_not_disclose_resolved_value( - self, - include_detailed_errors: bool, - ) -> None: - """Validation failures redact values resolved by security middleware.""" + assert observed == [] + assert executed is False + assert cast(dict[str, Any], context.arguments)["value"] == reference + + async def test_direct_composition_fails_before_allowed_hidden_execution(self) -> None: + """Allowed hidden arguments require a terminal middleware-pipeline boundary.""" tracker = LabelTrackingFunctionMiddleware() policy = PolicyEnforcementFunctionMiddleware() variable_id = tracker.get_variable_store().store( - "secret payload", + "payload", ContentLabel(integrity=IntegrityLabel.UNTRUSTED), ) + sink = self._sink(accepts_untrusted=True) + reference = f"[{variable_id}]" + context = FunctionInvocationContext(function=sink, arguments={"value": reference}) + executed = False - class StrictArgs(BaseModel): - value: int + async def execute() -> None: + nonlocal executed + executed = True - strict_tool = FunctionTool( - func=lambda value: str(value), - name="strict_sink", - input_model=StrictArgs, - additional_properties={"accepts_untrusted": True}, - ) - function_call = Content.from_function_call( - call_id="hidden-validation-error", + with pytest.raises(MiddlewareFailure, match="require FunctionMiddlewarePipeline"): + await tracker.process(context, lambda: policy.process(context, execute)) + + assert executed is False + assert cast(dict[str, Any], context.arguments)["value"] == reference + + async def test_equal_value_reference_substitution_fails_closed(self) -> None: + """Materialization binds the exact references and labels inspected by policy.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + safe_reference = tracker.get_variable_store().store( + "same payload", + ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=ConfidentialityLabel.PUBLIC, + ), + ) + restricted_reference = tracker.get_variable_store().store( + "same payload", + ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + metadata={ + _PRINCIPALS_KEY: [ + { + "tenant_id": "tenant-a", + "user_id": "user-a", + } + ] + }, + ), + ) + executed = False + + async def sink(value: str) -> str: + nonlocal executed + executed = True + return value + + sink_tool = FunctionTool( + func=sink, + name="public_reference_sink", + additional_properties={ + "accepts_untrusted": False, + "max_allowed_confidentiality": "public", + "source_integrity": "trusted", + }, + ) + + class SubstituteReferenceMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.arguments = {"value": f"[{restricted_reference}]"} + await call_next() + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await _auto_invoke_function( + Content.from_function_call( + call_id="equal-value-reference-substitution", + name=sink_tool.name, + arguments={"value": f"[{safe_reference}]"}, + ), + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + middleware_pipeline=FunctionMiddlewarePipeline( + tracker, + policy, + SubstituteReferenceMiddleware(), + ), + ) + + assert executed is False + + async def test_runtime_kwarg_mutation_after_approval_fails_closed(self) -> None: + """Approved runtime kwargs remain bound through final materialization.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink(accepts_untrusted=False) + request = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + kwargs={"tenant": "one"}, + ) + request.metadata["call_id"] = "runtime-kwarg-mutation" + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + request, + lambda _context: pytest.fail("Tool must wait for approval"), + ) + approval_request = cast(Content, request.result) + + class MutateRuntimeKwargMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.kwargs["tenant"] = "two" + await call_next() + + replay = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + kwargs={"tenant": "one"}, + ) + replay.metadata.update({ + "call_id": "runtime-kwarg-mutation", + "approval_response": approval_request.to_function_approval_response(True), + }) + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await FunctionMiddlewarePipeline(tracker, policy, MutateRuntimeKwargMiddleware()).execute( + replay, + lambda _context: pytest.fail("Mutated approved runtime kwargs must not execute"), + ) + + async def test_runtime_session_substitution_after_approval_fails_closed(self) -> None: + """Runtime AgentSession kwargs use the same normalized identity as approval binding.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink(accepts_untrusted=False) + request = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + kwargs={"session": AgentSession(session_id="tenant-one")}, + ) + request.metadata["call_id"] = "runtime-session-substitution" + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + request, + lambda _context: pytest.fail("Tool must wait for approval"), + ) + approval_request = cast(Content, request.result) + + class SubstituteRuntimeSessionMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.kwargs["session"] = AgentSession(session_id="tenant-two") + await call_next() + + replay = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + kwargs={"session": AgentSession(session_id="tenant-one")}, + ) + replay.metadata.update({ + "call_id": "runtime-session-substitution", + "approval_response": approval_request.to_function_approval_response(True), + }) + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await FunctionMiddlewarePipeline(tracker, policy, SubstituteRuntimeSessionMiddleware()).execute( + replay, + lambda _context: pytest.fail("Substituted runtime session must not execute"), + ) + + async def test_invocation_session_substitution_after_policy_fails_closed(self) -> None: + """The exact invocation session remains bound independently of runtime kwargs.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + sink = self._sink(accepts_untrusted=True) + original_session = AgentSession(session_id="shared-session-id") + context = FunctionInvocationContext( + function=sink, + arguments={"value": "payload"}, + session=original_session, + ) + executed = False + + class SubstituteInvocationSessionMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.session = AgentSession(session_id=original_session.session_id) + await call_next() + + async def execute(_: FunctionInvocationContext) -> None: + nonlocal executed + executed = True + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await FunctionMiddlewarePipeline(tracker, policy, SubstituteInvocationSessionMiddleware()).execute( + context, + execute, + ) + + assert executed is False + + async def test_security_scope_substitution_after_policy_fails_closed(self) -> None: + """The active FIDES authorization scope remains bound through execution.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + sink = self._sink(accepts_untrusted=True) + context = FunctionInvocationContext( + function=sink, + arguments={"value": "payload"}, + session=AgentSession(session_id="scope-owner"), + ) + replacement_scope = tracker._scope_for_session(AgentSession(session_id="other-scope")) + executed = False + + class SubstituteScopeMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + scope_token = tracker._active_security_scope.set(replacement_scope) + try: + await call_next() + finally: + tracker._active_security_scope.reset(scope_token) + + async def execute(_: FunctionInvocationContext) -> None: + nonlocal executed + executed = True + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await FunctionMiddlewarePipeline(tracker, policy, SubstituteScopeMiddleware()).execute(context, execute) + + assert executed is False + + @pytest.mark.parametrize("mutation", ["policy", "function"]) + async def test_function_policy_mutation_after_security_fails_closed(self, mutation: str) -> None: + """The exact tool object and parsed FIDES policy remain bound through execution.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + sink = self._sink(accepts_untrusted=True, max_confidentiality="private") + replacement = self._sink(accepts_untrusted=True, max_confidentiality="private") + context = FunctionInvocationContext(function=sink, arguments={"value": "payload"}) + executed = False + + class MutateFunctionPolicyMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + if mutation == "policy": + properties = cast(dict[str, Any], context.function.additional_properties) + properties["max_allowed_confidentiality"] = "public" + else: + context.function = replacement + await call_next() + + async def execute(_: FunctionInvocationContext) -> None: + nonlocal executed + executed = True + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await FunctionMiddlewarePipeline(tracker, policy, MutateFunctionPolicyMiddleware()).execute( + context, + execute, + ) + + assert executed is False + + async def test_destination_principal_mutation_after_security_fails_closed(self) -> None: + """The canonical USER_IDENTITY destination principal set remains bound through execution.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + tracker._context_label = ContentLabel( + confidentiality=ConfidentialityLabel.USER_IDENTITY, + metadata=_principal_metadata("user-a"), + ) + sink = _identity_destination(_principal_metadata("user-a")[_PRINCIPALS_KEY]) + context = FunctionInvocationContext(function=sink, arguments={"value": "payload"}) + executed = False + + class MutateDestinationPrincipalsMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + properties = cast(dict[str, Any], context.function.additional_properties) + properties[_PRINCIPALS_KEY] = _principal_metadata("user-b")[_PRINCIPALS_KEY] + await call_next() + + async def execute(_: FunctionInvocationContext) -> None: + nonlocal executed + executed = True + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await FunctionMiddlewarePipeline(tracker, policy, MutateDestinationPrincipalsMiddleware()).execute( + context, + execute, + ) + + assert executed is False + + async def test_runtime_bytes_mutation_after_policy_fails_closed(self) -> None: + """Immutable non-JSON runtime authority remains bound exactly.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + sink = self._sink(accepts_untrusted=True) + executed = False + + class MutateRuntimeCredentialMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.kwargs["credential"] = b"two" + await call_next() + + context = FunctionInvocationContext( + function=sink, + arguments={"value": "payload"}, + kwargs={"credential": b"one"}, + ) + + async def execute(_: FunctionInvocationContext) -> None: + nonlocal executed + executed = True + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await FunctionMiddlewarePipeline(tracker, policy, MutateRuntimeCredentialMiddleware()).execute( + context, + execute, + ) + + assert executed is False + + async def test_opaque_runtime_authority_fails_closed_before_downstream_middleware(self) -> None: + """Opaque application authority is rejected instead of omitted from binding.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + sink = self._sink(accepts_untrusted=True) + downstream_called = False + + class ObserveDownstreamMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + nonlocal downstream_called + downstream_called = True + await call_next() + + context = FunctionInvocationContext( + function=sink, + arguments={"value": "payload"}, + kwargs={"tenant": {"client": object()}}, + ) + + with pytest.raises(MiddlewareFailure, match="opaque runtime authority"): + await FunctionMiddlewarePipeline(tracker, policy, ObserveDownstreamMiddleware()).execute( + context, + lambda _context: pytest.fail("Opaque runtime authority must not execute"), + ) + + assert downstream_called is False + + async def test_direct_keyword_invocation_materializes_hidden_argument(self) -> None: + """FIDES preserves FunctionTool.invoke keyword arguments at the terminal boundary.""" + tracker = LabelTrackingFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + received: list[str] = [] + + async def sink(value: str) -> str: + received.append(value) + return value + + sink_tool = FunctionTool( + func=sink, + name="keyword_sink", + additional_properties={"accepts_untrusted": True}, + ) + reference = f"[{variable_id}]" + context = FunctionInvocationContext(function=sink_tool, arguments={"value": reference}) + + async def execute() -> None: + context.result = await sink_tool.invoke(value=reference) + + await tracker.process(context, execute) + + assert received == ["payload"] + assert cast(list[Content], context.result)[0].additional_properties["_variable_reference"] is True + + async def test_child_task_cannot_reuse_parent_materialization(self) -> None: + """A child task cannot inherit and replay a completed parent invocation.""" + tracker = LabelTrackingFunctionMiddleware() + received: list[str] = [] + release_child = asyncio.Event() + child_task: asyncio.Task[list[Content]] | None = None + + async def sink(value: str) -> str: + received.append(value) + return value + + sink_tool = FunctionTool( + func=sink, + name="task_local_sink", + additional_properties={"accepts_untrusted": True}, + ) + context = FunctionInvocationContext(function=sink_tool, arguments={"value": "parent"}) + + async def child_invoke() -> list[Content]: + await release_child.wait() + return await sink_tool.invoke(value="child") + + async def execute() -> None: + nonlocal child_task + child_task = asyncio.create_task(child_invoke()) + context.result = await sink_tool.invoke(value="parent") + + await tracker.process(context, execute) + assert child_task is not None + release_child.set() + child_result = await child_task + + assert received == ["parent", "child"] + assert child_result[0].text == "child" + + async def test_nested_same_task_invocation_cannot_reuse_parent_materialization(self) -> None: + """User tool code cannot reuse its parent's materialization capability.""" + tracker = LabelTrackingFunctionMiddleware() + received: list[str] = [] + nested_result: list[Content] = [] + + async def sink(value: str) -> str: + received.append(value) + if value == "parent": + nested_result.extend(await sink_tool.invoke(value="child")) + return value + + sink_tool = FunctionTool( + func=sink, + name="nested_task_local_sink", + additional_properties={"accepts_untrusted": True}, + ) + context = FunctionInvocationContext(function=sink_tool, arguments={"value": "parent"}) + + async def execute() -> None: + context.result = await sink_tool.invoke(value="parent") + + await tracker.process(context, execute) + + assert received == ["parent", "child"] + assert nested_result[0].text == "child" + + async def test_callable_policy_blocks_before_hidden_value_validation(self) -> None: + """A wrapped policy callable still authorizes before hidden-value validation.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + secret = "wrapped-policy-validator-secret" + variable_id = tracker.get_variable_store().store( + secret, + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + observed: list[str] = [] + executed = False + + class SinkArgs(BaseModel): + value: str + + @field_validator("value") + @classmethod + def observe_value(cls, value: str) -> str: + observed.append(value) + return value + + async def sink(value: str) -> str: + nonlocal executed + executed = True + return value + + sink_tool = FunctionTool( + func=sink, + name="wrapped_policy_sink", + input_model=SinkArgs, + additional_properties={ + "accepts_untrusted": False, + "max_allowed_confidentiality": "private", + "source_integrity": "trusted", + }, + ) + + with pytest.raises(MiddlewareTermination): + await _auto_invoke_function( + Content.from_function_call( + call_id="wrapped-policy-validator", + name=sink_tool.name, + arguments={"value": f"[{variable_id}]"}, + ), + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy.process), + ) + + assert secret not in observed + assert executed is False + + async def test_hidden_argument_retry_rematerializes_from_original_reference(self) -> None: + """A retrying outer middleware gets a fresh hidden-value materialization.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + secret = "retry-hidden-secret" + variable_id = tracker.get_variable_store().store( + secret, + ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ), + ) + attempts = 0 + received: list[str] = [] + + class SinkArgs(BaseModel): + value: str + + async def sink(value: str) -> str: + nonlocal attempts + attempts += 1 + received.append(value) + if attempts == 1: + raise RuntimeError("retry once") + return value + + class RetryOnceMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + del context + try: + await call_next() + except RuntimeError: + await call_next() + + sink_tool = FunctionTool( + func=sink, + name="retry_hidden_sink", + input_model=SinkArgs, + additional_properties={ + "accepts_untrusted": True, + "max_allowed_confidentiality": "private", + "source_integrity": "trusted", + }, + ) + + function_result = await _auto_invoke_function( + Content.from_function_call( + call_id="retry-hidden", + name=sink_tool.name, + arguments={"value": f"[{variable_id}]"}, + ), + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(RetryOnceMiddleware(), tracker, policy), + ) + + assert function_result.result == secret + assert attempts == 2 + assert received == [secret, secret] + + @pytest.mark.parametrize("mutation", ["tenant", "session", "context_session"]) + async def test_retry_after_policy_revalidates_materialized_runtime_authority(self, mutation: str) -> None: + """An inner retry cannot reuse materialization after changing runtime authority.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + attempts: list[tuple[str, str, str]] = [] + + async def sink(value: str, ctx: FunctionInvocationContext) -> str: + runtime_session = cast(AgentSession, ctx.kwargs["session"]) + invocation_session = cast(AgentSession, ctx.session) + attempts.append(( + cast(str, ctx.kwargs["tenant"]), + runtime_session.session_id, + invocation_session.session_id, + )) + if len(attempts) == 1: + raise RuntimeError("retry once") + return value + + class RetryAfterPolicy(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + try: + await call_next() + except RuntimeError: + if mutation == "tenant": + context.kwargs["tenant"] = "tenant-two" + elif mutation == "session": + context.kwargs["session"] = AgentSession(session_id="session-two") + else: + context.session = AgentSession(session_id="session-one") + await call_next() + + sink_tool = FunctionTool( + func=sink, + name="retry_authority_sink", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "private", + }, + ) + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await _auto_invoke_function( + Content.from_function_call( + call_id=f"retry-authority-{mutation}", + name=sink_tool.name, + arguments={"value": "payload"}, + ), + custom_args={"tenant": "tenant-one"}, + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + invocation_session=AgentSession(session_id="session-one"), + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy, RetryAfterPolicy()), + ) + + assert attempts == [("tenant-one", "session-one", "session-one")] + + @pytest.mark.parametrize("value", ["[var_0123456789abcdef]", "var_0123456789abcdef"]) + async def test_unknown_references_remain_literal(self, value: str) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + received: list[str] = [] + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=False), + arguments={"value": value}, + ) + + async def execute(current: FunctionInvocationContext) -> list[Content]: + received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + + assert received == [value] + argument_label = cast(ContentLabel, context.metadata["argument_label"]) + assert argument_label.integrity == IntegrityLabel.TRUSTED + assert argument_label.confidentiality == ConfidentialityLabel.PUBLIC + + async def test_accepts_untrusted_blind_forwards_without_model_context_taint(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + received: list[str] = [] + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=True), + arguments={"value": f"[{variable_id}]"}, + ) + + async def execute(current: FunctionInvocationContext) -> list[Content]: + received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + + assert received == ["payload"] + assert context.metadata["argument_label"].integrity == IntegrityLabel.UNTRUSTED + assert tracker.get_context_label().integrity == IntegrityLabel.TRUSTED + + async def test_hidden_argument_resolution_does_not_require_reapproval(self) -> None: + """Security expansion preserves the approval-visible placeholder.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink(accepts_untrusted=True) + function_call = Content.from_function_call( + call_id="approved-hidden-value", + id="approved-hidden-value-occurrence", + name=sink.name, + arguments={"value": f"[{variable_id}]"}, + ) + approval_response = Content.from_function_approval_request( + id="approved-hidden-value-occurrence", + function_call=function_call, + ).to_function_approval_response(approved=True) + + result = await _auto_invoke_function( + approval_response, + config=normalize_function_invocation_configuration(None), + tool_map={sink.name: sink}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) + + assert result.type == "function_result" + assert result.result == "payload" + + async def test_hidden_argument_can_be_normalized_after_security_check(self) -> None: + """Final Pydantic coercion is not mistaken for post-policy mutation.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "3", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + + class StrictArgs(BaseModel): + value: int + + received: list[int] = [] + + def strict_sink(value: int) -> str: + received.append(value) + return str(value) + + strict_tool = FunctionTool( + func=strict_sink, + name="strict_sink", + input_model=StrictArgs, + additional_properties={"accepts_untrusted": True}, + ) + function_call = Content.from_function_call( + call_id="hidden-normalization", + name=strict_tool.name, + arguments={"value": f"[{variable_id}]"}, + ) + + result = await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={strict_tool.name: strict_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) + + assert result.type == "function_result" + assert result.exception is None + assert received == [3] + + @pytest.mark.parametrize("include_detailed_errors", [False, True]) + async def test_hidden_argument_validation_error_does_not_disclose_resolved_value( + self, + include_detailed_errors: bool, + ) -> None: + """Validation failures redact values resolved by security middleware.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "secret payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + + class StrictArgs(BaseModel): + value: int + + strict_tool = FunctionTool( + func=lambda value: str(value), + name="strict_sink", + input_model=StrictArgs, + additional_properties={"accepts_untrusted": True}, + ) + function_call = Content.from_function_call( + call_id="hidden-validation-error", name=strict_tool.name, arguments={"value": f"[{variable_id}]"}, ) @@ -7943,6 +8856,62 @@ def transformed_sink(value: str) -> str: assert received == ["dangerous-operation"] assert validation_count == 1 + async def test_policy_approval_rejects_validator_transformed_hidden_argument(self) -> None: + """A FIDES policy grant cannot authorize a different validator-normalized invocation.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + variable_id = tracker.get_variable_store().store( + "safe", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + executed = False + + class TransformArgs(BaseModel): + value: str + + @field_validator("value") + @classmethod + def transform_value(cls, value: str) -> str: + return "dangerous-operation" if value == "safe" else value + + async def sink(value: str) -> str: + nonlocal executed + executed = True + return value + + sink_tool = FunctionTool( + func=sink, + name="approved_transform_sink", + input_model=TransformArgs, + additional_properties={"accepts_untrusted": False}, + ) + function_call = Content.from_function_call( + call_id="approved-validator-transform", + id="approved-validator-transform-occurrence", + name=sink_tool.name, + arguments={"value": f"[{variable_id}]"}, + ) + pipeline = FunctionMiddlewarePipeline(tracker, policy) + + with pytest.raises(MiddlewareTermination) as approval_exc: + await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + middleware_pipeline=pipeline, + ) + approval_request = cast(Content, approval_exc.value.result) + + with pytest.raises(MiddlewareFailure, match="changed during validation after FIDES policy approval"): + await _auto_invoke_function( + approval_request.to_function_approval_response(True), + config=normalize_function_invocation_configuration(None), + tool_map={sink_tool.name: sink_tool}, + middleware_pipeline=pipeline, + ) + + assert executed is False + async def test_security_policy_binds_values_converted_by_input_model(self) -> None: """Security middleware can snapshot datetime and set arguments converted from JSON (#8661).""" tracker = LabelTrackingFunctionMiddleware() @@ -8246,15 +9215,12 @@ async def test_opaque_runtime_kwarg_fails_closed(self) -> None: ) context.metadata["call_id"] = "call-opaque" - with pytest.raises(MiddlewareTermination): + with pytest.raises(MiddlewareFailure, match="opaque runtime authority"): await FunctionMiddlewarePipeline(tracker, policy).execute( context, lambda _context: pytest.fail("Opaque invocation must not execute"), ) - assert isinstance(context.result, dict) - assert context.result["violation_type"] == "unsafe_approval_binding" - async def test_pending_policy_approval_survives_session_restore(self) -> None: config = SecureAgentConfig(approval_on_violation=True) session = AgentSession(session_id="restored-approval") @@ -8701,10 +9667,6 @@ async def my_tool(files: ReorderedFiles): ) async def call_next(): - context.function._prepare_context_arguments( # pyright: ignore[reportPrivateUsage] - context, - context.arguments, - ) await tool.invoke(arguments=context.arguments) await tracker.process(context, call_next) @@ -8733,10 +9695,6 @@ async def my_tool(files: FilteredFiles): ) async def call_next(): - context.function._prepare_context_arguments( # pyright: ignore[reportPrivateUsage] - context, - context.arguments, - ) await tool.invoke(arguments=context.arguments) await tracker.process(context, call_next) @@ -8765,10 +9723,6 @@ async def my_tool(files: PreservedFiles): ) async def call_next(): - context.function._prepare_context_arguments( # pyright: ignore[reportPrivateUsage] - context, - context.arguments, - ) await tool.invoke(arguments=context.arguments) await tracker.process(context, call_next) diff --git a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md index b574ddd115..0d10caaa86 100644 --- a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md +++ b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md @@ -1094,10 +1094,18 @@ Configure tool security requirements in the `@tool` decorator: ``` Hidden variable references in tool arguments are expanded recursively, but their stored labels remain attached to -the invocation. A tool with `accepts_untrusted=False` is blocked, audited, or sent for policy approval before it can -receive hidden untrusted data. `accepts_untrusted=True` permits blind forwarding without exposing that data to the -model context. It does not bypass `max_allowed_confidentiality`; hidden private data still cannot flow to a public -sink. Argument labels do not rewrite result labels. +the invocation. FIDES resolves a detached copy for label and approval binding, then blocks, audits, or requests policy +approval before it mutates the live arguments or runs Pydantic validators. Validators and the tool body receive hidden +values only after the policy permits the call. `accepts_untrusted=True` permits blind forwarding without exposing that +data to the model context. It does not bypass `max_allowed_confidentiality`; hidden private data still cannot flow to +a public sink. Argument labels do not rewrite result labels. After a policy violation is approved, validator +normalization must preserve the approved resolved invocation; a changed value fails closed before the tool body runs. + +FIDES metadata keys and values are exact and case-sensitive. Misspelled or incorrectly cased FIDES-like keys, +unknown label values, and non-boolean `accepts_untrusted` declarations are blocked as an `invalid_policy` +configuration violation instead of falling back to a less restrictive policy. Unrelated integration metadata remains +allowed. FIDES binds the exact tool object, normalized policy, and canonical USER_IDENTITY principal set through +terminal execution and retries, so later middleware cannot replace the authorization policy after preflight. **Approval model:** - Use `approval_mode="always_require"` for normal human-in-the-loop approval on a specific tool.