From eef825f19db4b769305c760e6de3aa2707e64a2f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 11 Sep 2026 11:43:49 +0200 Subject: [PATCH 1/4] Python: allow middleware to repair function arguments Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../specs/004-python-function-calling-loop.md | 16 + .../core/agent_framework/_middleware.py | 30 +- .../packages/core/agent_framework/_tools.py | 307 ++++++++++++++---- .../packages/core/agent_framework/security.py | 8 +- .../core/test_function_invocation_logic.py | 279 ++++++++++++++++ python/packages/core/tests/test_security.py | 164 +++++++++- python/samples/02-agents/tools/README.md | 8 +- .../function_tool_with_explicit_schema.py | 13 +- 8 files changed, 746 insertions(+), 79 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 81019bb5bbe..426a78ed959 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -377,6 +377,15 @@ that manually replay messages own the equivalent rule: do not resend an approval not invent one. - A completed function call/result pair is inert on later turns. - Informational-only and declaration-only calls are not executed as local tools. +- For automatic local execution, provider arguments are JSON-parsed before function middleware but are not schema + validated yet. Middleware may inspect or repair that raw mapping before calling `call_next()`. The innermost handler + performs final validation immediately before the tool body and writes the validated, normalized mapping back to + `FunctionInvocationContext.arguments`. Middleware that short-circuits without `call_next()` also skips final + validation and tool execution. +- Argument-repair middleware must precede security or policy middleware so enforcement observes the effective + invocation. Changing arguments after security middleware has processed them fails closed with `MiddlewareFailure`. +- Argument-validation failures after middleware retain the established `Argument parsing failed` result contract. + Exceptions raised by middleware or the tool body retain the separate `Function failed` contract. ### Reasoning-bound calls @@ -414,6 +423,11 @@ that manually replay messages own the equivalent rule: do not resend an approval - If policy middleware detects that the exact resolved invocation changed after approval, the old response executes nothing and yields a caller-visible, session-persisted replacement request for the same occurrence; execution requires a second approval and happens exactly once. +- If function middleware repairs an approval-bound call, the old response likewise executes nothing and yields a + caller-visible, session-persisted replacement request containing the repaired approval-visible arguments. The + replacement retains the call occurrence identity, rotates request-generation identity, and requires a second + approval before exactly-once execution. Security middleware transformations that preserve an approval-visible + placeholder do not disclose the resolved value or trigger a spurious replacement. - If session-bound middleware no longer holds the reviewed authority because it expired or was evicted, the matched response executes nothing and produces a replacement approval request with the same occurrence identity and a fresh request generation. The replacement is caller-visible, becomes the authoritative pending session snapshot, and @@ -583,6 +597,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Rejected execution | Rejection is a normal terminal result, not an exception to the caller. | `test_unapproved_tool_execution_raises_exception` | | Approved tool exception | Generic and detailed error modes preserve one result and one execution. | `test_approved_function_call_with_error_without_detailed_errors`, `test_approved_function_call_with_error_with_detailed_errors` | | Approved validation error | Validation failure returns one result without invoking the function body. | `test_approved_function_call_with_validation_error` | +| Pre-validation middleware repair | Function middleware observes raw parsed arguments, may repair them before final validation, and the body receives validated normalized values exactly once; short-circuiting skips validation and execution. Repair after security middleware fails closed, including invalid and short-circuited mutations. Validation errors after hidden-value resolution do not disclose the resolved value, while ordinary final Pydantic normalization remains allowed. | `test_function_middleware_repairs_raw_arguments_before_validation`, `test_function_middleware_can_short_circuit_before_argument_validation`, `test_invalid_arguments_produced_by_middleware_keep_argument_error_contract`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_argument_mutation_after_security_middleware_fails_closed`, `test_argument_mutation_after_security_short_circuit_fails_closed`, `test_hidden_argument_validation_error_does_not_disclose_resolved_value`, `test_hidden_argument_can_be_normalized_after_security_check` | +| Approved middleware repair | A changed approval-bound call executes zero times under the old grant, returns a persisted occurrence-bound replacement request in both response modes, and executes once only after the replacement is approved. Security expansion preserves approval-visible placeholders. | `test_approved_argument_repair_requires_replacement_approval`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_hidden_argument_resolution_does_not_require_reapproval` | | Approved success | Successful approved execution returns one result. | `test_approved_function_call_successful_execution` | | Consecutive error cap | Error threshold stops repeated failures, submits collected results, and makes only the required final no-tool model call. | `test_function_invocation_config_max_consecutive_errors`, `test_streaming_function_invocation_config_max_consecutive_errors`, `test_approval_resume_error_limit_forces_final_no_tool_response` | | Unknown call handling | Configured false returns an error result; configured true raises. | `test_function_invocation_config_terminate_on_unknown_calls_false`, `test_function_invocation_config_terminate_on_unknown_calls_true`, streaming equivalents | diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 419cff8f185..1423ec83b76 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -370,7 +370,13 @@ class FunctionInvocationContext: Attributes: function: The function being invoked. - arguments: The validated arguments for the function. + arguments: The function arguments. In the automatic function-calling loop, + middleware initially receives the raw JSON-parsed mapping from the + provider and may repair it before calling ``call_next()``. The innermost + handler validates and normalizes the current value immediately before + execution, then stores the normalized mapping back on this attribute. + Middleware that short-circuits without calling ``call_next()`` also skips + final schema validation and function execution. session: The agent session for this invocation, if any. metadata: Metadata dictionary for sharing data between function middleware. result: Function execution result. This attribute carries no guaranteed type. @@ -438,7 +444,9 @@ def __init__( Args: function: The function being invoked. - arguments: The validated arguments for the function. + arguments: The function arguments. Automatic invocation supplies the raw + JSON-parsed mapping; final validation occurs after middleware, immediately + before function execution. session: The agent session for this invocation, if any. metadata: Metadata dictionary for sharing data between function middleware. result: Function execution result. Observed and overridden values do not @@ -714,8 +722,15 @@ class FunctionMiddleware(ABC): """Abstract base class for function middleware that can intercept function invocations. Function middleware allows you to intercept and modify function/tool invocations before - and after execution. You can validate arguments, cache results, log invocations, or - override function execution. + and after execution. On entry, automatic function invocation exposes the raw JSON-parsed + arguments so middleware can repair provider-specific deviations before calling + ``call_next()``. The innermost handler performs final validation immediately before + execution and updates ``context.arguments`` with the normalized values. You can also + cache results, log invocations, or override function execution. + + Argument-repair middleware must run before security or policy middleware so those + layers inspect the effective invocation. Changing arguments after security middleware + has processed them fails closed with :class:`MiddlewareFailure`. Note: FunctionMiddleware is an abstract base class. You must subclass it and implement @@ -768,8 +783,11 @@ async def process( Args: context: Function invocation context containing function, arguments, and metadata. - MiddlewareTypes can set context.result to override execution, or observe - the actual execution result after calling call_next(). + Before ``call_next()``, automatic invocation exposes raw JSON-parsed + arguments that middleware may inspect or replace. After ``call_next()`` + reaches the function, arguments contain their validated, normalized + values. MiddlewareTypes can set context.result to override execution, + or observe the actual execution result after calling call_next(). call_next: Function to call the next middleware or final function execution. Does not return anything - all data flows through the context. diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index d09078ea434..c1299bab358 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -121,10 +121,30 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None) _FUNCTION_RESULT_CARRIER_CONTEXT_KEY: Final[str] = "_function_result_carrier" _FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY: Final[str] = "_function_result_payload_budget" _FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY: Final[str] = "_function_result_payload_budget" +_APPROVED_ARGUMENTS_CONTEXT_KEY: Final[str] = "_approved_function_arguments" +_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY: Final[str] = "_security_function_arguments" _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"} + + +class _FunctionArgumentValidationError(TypeError): + """An argument-validation failure raised before the function body starts.""" + + def __init__(self, message: str, *, redacted_message: str | None = None) -> None: + super().__init__(message) + self.redacted_message = redacted_message or message + + +class _FunctionArgumentsChangedAfterApproval(Exception): + """Signal that middleware changed an approval-bound invocation.""" + + def __init__(self, arguments: Mapping[str, Any]) -> None: + super().__init__("Function arguments changed after approval.") + self.arguments = copy.deepcopy(dict(arguments)) + + ApprovalMode: TypeAlias = Literal["always_require", "never_require"] ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]") ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) @@ -323,7 +343,10 @@ class FunctionTool(SerializationMixin): """A tool that wraps a Python function to make it callable by AI models. This class wraps a Python function to make it callable by AI models with automatic - parameter validation and JSON schema generation. + parameter validation and JSON schema generation. Inferred and Pydantic input models + provide recursive runtime validation. Caller-supplied JSON schema mappings are passed + through to providers and receive only the lightweight checks documented on + :paramref:`input_model`; they are not an authorization or security boundary. Attributes: name: The name of the tool. @@ -434,6 +457,15 @@ def __init__( parameters, explicitly provide ``input_model`` (either a Pydantic ``BaseModel`` or a JSON schema dictionary) so the model can reason about the expected arguments. + + A dictionary is preserved as supplied and checked only for top-level + ``required`` fields, ``additionalProperties: false``, and property + ``enum`` and primitive ``type`` values. Nested constraints, + compositions such as ``oneOf``, references such as ``$ref``, and other + JSON Schema keywords are not comprehensively enforced at runtime. Use a + Pydantic model when runtime validation matters. Treat dictionary schemas + as declarations for trusted settings and non-sensitive functions only; + never rely on them as an authorization or security boundary. result_parser: An optional callable with signature ``Callable[[Any], str]`` that overrides the default result parsing behavior. When provided, this callable is used to convert the raw function return value to a string instead of the @@ -631,6 +663,116 @@ async def _invoke_function(self, call_kwargs: Mapping[str, Any]) -> Any: res = await asyncio.to_thread(self.__call__, **call_kwargs) return await res if inspect.isawaitable(res) else res + def _prepare_arguments(self, arguments: BaseModel | Mapping[str, Any] | None) -> dict[str, Any]: + """Validate and normalize arguments immediately before function execution.""" + if arguments is None: + return {} + + try: + if isinstance(arguments, Mapping): + parsed_arguments = dict(arguments) + if self.input_model is not None and not self._schema_supplied: + # exclude_unset (not exclude_none): keep arguments the model + # explicitly provided even when their value is null, and drop + # only the ones it left out, so the function's own defaults + # apply. Excluding null instead would strip a required nullable + # parameter the model deliberately set to null, failing the + # invocation on the missing argument (#5934). + parsed_arguments = self.input_model.model_validate(parsed_arguments).model_dump(exclude_unset=True) + elif isinstance(arguments, BaseModel): + if ( + self.input_model is not None + and not self._schema_supplied + and not isinstance(arguments, self.input_model) + ): + raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}") + parsed_arguments = arguments.model_dump(exclude_unset=True) + else: + raise TypeError( + f"Expected mapping-like arguments for tool '{self.name}', got {type(arguments).__name__}" + ) + except ValidationError as exc: + locations = sorted({ + ".".join(str(part) for part in error["loc"]) + for error in exc.errors(include_input=False, include_url=False) + }) + redacted_message = f"Invalid arguments for '{self.name}'." + if locations: + redacted_message = f"{redacted_message} Invalid field(s): {', '.join(locations)}." + raise _FunctionArgumentValidationError( + f"Invalid arguments for '{self.name}': {exc}", + redacted_message=redacted_message, + ) from exc + except TypeError as exc: + raise _FunctionArgumentValidationError(str(exc)) from exc + + try: + return _validate_arguments_against_schema( + arguments=parsed_arguments, + schema=self.parameters(), + tool_name=self.name, + ) + except TypeError as exc: + raise _FunctionArgumentValidationError(str(exc)) from exc + + @staticmethod + def _arguments_as_mapping(arguments: Any) -> dict[str, Any] | None: + """Return arguments as a mapping without applying schema validation.""" + candidate = arguments + if candidate is None: + return {} + if isinstance(candidate, BaseModel): + return candidate.model_dump(exclude_unset=True) + if isinstance(candidate, Mapping): + return dict(cast(Mapping[str, Any], candidate)) + return None + + @classmethod + def _approval_visible_arguments( + cls, + arguments: BaseModel | Mapping[str, Any] | None, + context: FunctionInvocationContext | None, + ) -> dict[str, Any] | None: + """Return the non-expanded arguments that an approval request may disclose.""" + if context is not None and "original_arguments_for_messages" in context.metadata: + return cls._arguments_as_mapping(context.metadata["original_arguments_for_messages"]) + return cls._arguments_as_mapping(arguments) + + @staticmethod + def _ensure_security_arguments_unchanged( + context: FunctionInvocationContext | None, + current_arguments: Mapping[str, Any] | None, + ) -> None: + """Fail closed when arguments change after security middleware processed them.""" + if context is None: + return + security_arguments = context.metadata.get(_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY) + if isinstance(security_arguments, Mapping) and ( + current_arguments is None or dict(cast(Mapping[str, Any], security_arguments)) != dict(current_arguments) + ): + from ._middleware import MiddlewareFailure + + raise MiddlewareFailure( + "Function arguments changed after security middleware processed them. " + "Install argument-repair middleware before security middleware." + ) + + @staticmethod + def _ensure_approved_arguments_unchanged( + context: FunctionInvocationContext | None, + approval_visible_arguments: Mapping[str, Any] | None, + ) -> None: + """Require a replacement approval when middleware changes approved arguments.""" + if context is None: + return + if approval_visible_arguments is None: + return + approved_arguments = context.metadata.get(_APPROVED_ARGUMENTS_CONTEXT_KEY) + if not isinstance(approved_arguments, Mapping): + return + if dict(cast(Mapping[str, Any], approved_arguments)) != dict(approval_visible_arguments): + raise _FunctionArgumentsChangedAfterApproval(approval_visible_arguments) + @overload async def invoke( self, @@ -722,42 +864,10 @@ async def invoke( if arguments is None and context is not None: arguments = context.arguments - if arguments is None: - validated_arguments: dict[str, Any] = {} - else: - try: - if isinstance(arguments, Mapping): - parsed_arguments = dict(arguments) - if self.input_model is not None and not self._schema_supplied: - # exclude_unset (not exclude_none): keep arguments the model - # explicitly provided even when their value is null, and drop - # only the ones it left out, so the function's own defaults - # apply. Excluding null instead would strip a required nullable - # parameter the model deliberately set to null, failing the - # invocation on the missing argument (#5934). - parsed_arguments = self.input_model.model_validate(parsed_arguments).model_dump( - exclude_unset=True - ) - elif isinstance(arguments, BaseModel): - if ( - self.input_model is not None - and not self._schema_supplied - and not isinstance(arguments, self.input_model) - ): - raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}") - parsed_arguments = arguments.model_dump(exclude_unset=True) - else: - raise TypeError( - f"Expected mapping-like arguments for tool '{self.name}', got {type(arguments).__name__}" - ) - except ValidationError as exc: - raise TypeError(f"Invalid arguments for '{self.name}': {exc}") from exc - - validated_arguments = _validate_arguments_against_schema( - arguments=parsed_arguments, - schema=self.parameters(), - tool_name=self.name, - ) + current_arguments = self._arguments_as_mapping(arguments) + approval_visible_arguments = self._approval_visible_arguments(arguments, context) + self._ensure_security_arguments_unchanged(context, current_arguments) + validated_arguments = self._prepare_arguments(arguments) effective_context = context if effective_context is None and self._context_parameter_name is not None: @@ -771,6 +881,8 @@ async def invoke( effective_context.arguments = validated_arguments effective_context.kwargs = dict(runtime_kwargs) + self._ensure_approved_arguments_unchanged(effective_context, approval_visible_arguments) + call_kwargs = dict(validated_arguments) observable_kwargs = dict(validated_arguments) if self._context_parameter_name is not None and effective_context is not None: @@ -1177,7 +1289,7 @@ def _validate_arguments_against_schema( schema: Mapping[str, Any], tool_name: str, ) -> dict[str, Any]: - """Run lightweight argument checks for schema-supplied tools.""" + """Run lightweight, top-level argument checks for schema-supplied tools.""" parsed_arguments = dict(arguments) required_fields = [field for field in schema.get("required", []) if isinstance(field, str)] @@ -1197,9 +1309,7 @@ def _validate_arguments_against_schema( enum_values = properties.get(field_name, {}).get("enum") if isinstance(enum_values, list) and enum_values and field_value not in enum_values: - raise TypeError( - f"Invalid value for '{field_name}' in '{tool_name}': {field_value!r} is not in {enum_values!r}" - ) + raise TypeError(f"Invalid value for '{field_name}' in '{tool_name}': value is not in {enum_values!r}") schema_type = properties.get(field_name, {}).get("type") if isinstance(schema_type, str): @@ -1293,9 +1403,14 @@ def tool( docstring will be used. schema: An explicit input schema for the function. This can be a Pydantic ``BaseModel`` subclass or a JSON schema dictionary (``Mapping[str, Any]``). - When a dictionary is provided, it must be a flat object schema with a - ``properties`` key (complex JSON Schema features such as ``oneOf``, - ``$ref``, or nested compositions are not supported). + Dictionary schemas are passed through to providers and receive only + lightweight top-level checks for ``required``, ``additionalProperties: + false``, property ``enum``, and primitive property ``type``. Nested + constraints, compositions such as ``oneOf``, references such as ``$ref``, + and other JSON Schema keywords are not comprehensively enforced at runtime. + Use a Pydantic model when runtime validation matters. Dictionary schemas are + intended for trusted settings and non-sensitive functions and must not be + treated as an authorization or security boundary. When provided, the schema is used instead of inferring one from the function's signature. Defaults to ``None`` (infer from signature). approval_mode: Whether or not approval is required to run this tool. @@ -1542,6 +1657,62 @@ def _function_execution_error_result( ) +def _function_argument_validation_error_result( + function_call: Content, + exception: _FunctionArgumentValidationError, + config: FunctionInvocationConfiguration, + context: FunctionInvocationContext | None = None, +) -> Content: + """Build the stable tool result for argument-validation failures.""" + from ._types import Content + + exception_message = ( + exception.redacted_message + if context is not None and _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY in context.metadata + else str(exception) + ) + message = "Error: Argument parsing failed." + if config.get("include_detailed_errors", False): + message = f"{message} Exception: {exception_message}" + return Content.from_function_result( + call_id=function_call.call_id, # type: ignore[arg-type] + result=message, + exception=exception_message, + additional_properties=function_call.additional_properties, + ) + + +def _replacement_approval_request( + function_call: Content, + arguments: Mapping[str, Any], +) -> Content: + """Create a new approval generation for middleware-repaired arguments.""" + from ._types import Content + + call_id = function_call.call_id + if call_id is None: + raise KeyError(f'Function "{function_call.name}" is missing call_id.') + occurrence_id = function_call.id or call_id + request_id = f"{occurrence_id}:replacement:{uuid4().hex}" + repaired_call = Content.from_function_call( + call_id=call_id, + name=function_call.name, # type: ignore[arg-type] + arguments=copy.deepcopy(dict(arguments)), + id=occurrence_id, + annotations=copy.deepcopy(function_call.annotations), + additional_properties=copy.deepcopy(function_call.additional_properties), + ) + return Content.from_function_approval_request( + id=request_id, + function_call=repaired_call, + additional_properties={ + _APPROVAL_REQUEST_ID_KEY: request_id, + "_replacement_approval_request": True, + "reason": "Function arguments changed after approval.", + }, + ) + + def _finalize_function_result( *, call_id: str, @@ -1672,29 +1843,7 @@ async def _auto_invoke_function( } if invocation_session is not None: runtime_kwargs["session"] = invocation_session - try: - if not cast(bool, getattr(tool, "_schema_supplied", False)) and tool.input_model is not None: - # exclude_unset (not exclude_none) so an argument the model explicitly set - # to null still reaches the function; see FunctionTool.invoke for the full - # rationale. This is the auto-calling path #5934 actually hits. - args = tool.input_model.model_validate(parsed_args).model_dump(exclude_unset=True) - else: - args = dict(parsed_args) - args = _validate_arguments_against_schema( - arguments=args, - schema=tool.parameters(), - tool_name=tool.name, - ) - except (TypeError, ValidationError) as exc: - message = "Error: Argument parsing failed." - if config.get("include_detailed_errors", False): - message = f"{message} Exception: {exc}" - return Content.from_function_result( - call_id=function_call_content.call_id, # type: ignore[arg-type] - result=message, - exception=str(exc), - additional_properties=function_call_content.additional_properties, - ) + args = dict(parsed_args) from ._middleware import FunctionInvocationContext, MiddlewareFailure @@ -1727,6 +1876,8 @@ async def _auto_invoke_function( # Explicit control-flow signals escape the loop; only ordinary exceptions # are absorbed into tool-error results below. raise + except _FunctionArgumentValidationError as exc: + return _function_argument_validation_error_result(function_call_content, exc, config) except Exception as exc: return _function_execution_error_result(function_call_content, tool.name, exc, config, direct_context) # Execute through middleware pipeline if available @@ -1753,8 +1904,13 @@ async def _auto_invoke_function( # this replay corresponds to a middleware-specific approval flow. if approval_response is not None: middleware_context.metadata["approval_response"] = approval_response + middleware_context.metadata[_APPROVED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(parsed_args) + + final_handler_started = False async def final_function_handler(context_obj: Any) -> Any: + nonlocal final_handler_started + final_handler_started = True return await tool.invoke( arguments=context_obj.arguments, context=context_obj, @@ -1769,6 +1925,11 @@ async def final_function_handler(context_obj: Any) -> Any: context=middleware_context, final_handler=final_function_handler, ) + if not final_handler_started: + tool._ensure_security_arguments_unchanged( # pyright: ignore[reportPrivateUsage] + middleware_context, + tool._arguments_as_mapping(middleware_context.arguments), # pyright: ignore[reportPrivateUsage] + ) # Pass through function_approval_request directly (e.g., from security middleware) if isinstance(function_result, Content) and function_result.type == "function_approval_request": @@ -1781,6 +1942,11 @@ async def final_function_handler(context_obj: Any) -> Any: context=middleware_context, ) except MiddlewareTermination as term_exc: + if not final_handler_started: + tool._ensure_security_arguments_unchanged( # pyright: ignore[reportPrivateUsage] + middleware_context, + tool._arguments_as_mapping(middleware_context.arguments), # pyright: ignore[reportPrivateUsage] + ) # Re-raise to signal loop termination, but first capture any result set by middleware if middleware_context.result is not None: # Pass through function_approval_request directly (e.g., from security policy middleware) @@ -1799,6 +1965,13 @@ async def final_function_handler(context_obj: Any) -> Any: context=middleware_context, ) raise + except _FunctionArgumentsChangedAfterApproval as exc: + raise MiddlewareTermination( + "Function arguments changed after approval.", + result=_replacement_approval_request(function_call_content, exc.arguments), + ) from exc + except _FunctionArgumentValidationError as exc: + return _function_argument_validation_error_result(function_call_content, exc, config, middleware_context) except (MiddlewareFailure, UserInputRequiredException): # MiddlewareFailure is the loop's explicit fail-closed escape: middleware that # must abort the run (enforcement layers, guardrails) raises it instead of diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index a655a9971af..c4acbf33d3f 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -36,7 +36,12 @@ from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination from ._serialization import SerializationMixin from ._sessions import AgentSession, ContextProvider -from ._tools import _APPROVAL_REQUEST_ID_KEY, FunctionTool, tool # pyright: ignore[reportPrivateUsage] +from ._tools import ( + _APPROVAL_REQUEST_ID_KEY, # pyright: ignore[reportPrivateUsage] + _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] + FunctionTool, + tool, +) from ._types import Content, Message if TYPE_CHECKING: @@ -1549,6 +1554,7 @@ async def process( # 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] = deepcopy(context.arguments) argument_labels = [*input_labels, *resolved_labels] argument_label = combine_labels(*argument_labels) if argument_labels else ContentLabel() diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 7e688a2e13c..80a50de2d6a 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3122,6 +3122,285 @@ def typed_func(arg1: int) -> str: # Expects int, not str assert "Exception:" not in error_result.result # No detailed error +async def test_function_middleware_repairs_raw_arguments_before_validation( + chat_client_base: SupportsChatGetResponse, +) -> None: + """Function middleware can repair provider arguments before final validation.""" + observed_arguments: list[dict[str, Any]] = [] + executed_arguments: list[int] = [] + + class RepairArgumentsMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + assert isinstance(context.arguments, dict) + observed_arguments.append(dict(context.arguments)) + context.arguments["count"] = int(context.arguments.pop("count_text")) + await call_next() + assert context.arguments == {"count": 3} + + @tool(name="count_items", approval_mode="never_require") + def count_items(count: int) -> str: + executed_arguments.append(count) + return str(count) + + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="repair-1", + name="count_items", + arguments='{"count_text": "3"}', + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + await chat_client_base.get_response( + [Message(role="user", contents=["count"])], + options={"tools": [count_items]}, + client_kwargs={"middleware": [RepairArgumentsMiddleware()]}, + ) + + assert observed_arguments == [{"count_text": "3"}] + assert executed_arguments == [3] + + +async def test_function_middleware_can_short_circuit_before_argument_validation( + chat_client_base: SupportsChatGetResponse, +) -> None: + """A middleware result can handle an invalid call without invoking validation or the tool.""" + executions = 0 + + class ShortCircuitMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + del call_next + assert context.arguments == {"wrong_name": "not-an-int"} + context.result = "handled by middleware" + + @tool(name="strict_tool", approval_mode="never_require") + def strict_tool(count: int) -> str: + nonlocal executions + executions += 1 + return str(count) + + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="short-circuit-1", + name="strict_tool", + arguments='{"wrong_name": "not-an-int"}', + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + response = await chat_client_base.get_response( + [Message(role="user", contents=["count"])], + options={"tools": [strict_tool]}, + client_kwargs={"middleware": [ShortCircuitMiddleware()]}, + ) + + assert executions == 0 + result = next( + content for message in response.messages for content in message.contents if content.type == "function_result" + ) + assert result.result == "handled by middleware" + + +async def test_invalid_arguments_produced_by_middleware_keep_argument_error_contract( + chat_client_base: SupportsChatGetResponse, +) -> None: + """Final validation still rejects invalid middleware output before the tool body.""" + executions = 0 + + class InvalidRepairMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.arguments = {"count": "not-an-int"} + await call_next() + + @tool(name="strict_tool", approval_mode="never_require") + def strict_tool(count: int) -> str: + nonlocal executions + executions += 1 + return str(count) + + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="invalid-repair-1", name="strict_tool", arguments='{"count": 1}') + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + response = await chat_client_base.get_response( + [Message(role="user", contents=["count"])], + options={"tools": [strict_tool]}, + client_kwargs={"middleware": [InvalidRepairMiddleware()]}, + ) + + assert executions == 0 + result = next( + content for message in response.messages for content in message.contents if content.type == "function_result" + ) + assert result.result == "Error: Argument parsing failed." + assert result.exception is not None + + +async def test_middleware_type_error_remains_function_error(chat_client_base: SupportsChatGetResponse) -> None: + """A middleware TypeError is not misclassified as an argument-validation failure.""" + + class FailingMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + del context, call_next + raise TypeError("middleware failed") + + @tool(name="strict_tool", approval_mode="never_require") + def strict_tool(count: int) -> str: + return str(count) + + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="middleware-error-1", name="strict_tool", arguments='{"count": 1}' + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + response = await chat_client_base.get_response( + [Message(role="user", contents=["count"])], + options={"tools": [strict_tool]}, + client_kwargs={"middleware": [FailingMiddleware()]}, + ) + + result = next( + content for message in response.messages for content in message.contents if content.type == "function_result" + ) + assert result.result == "Error: Function failed." + assert result.exception == "middleware failed" + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_approved_argument_repair_requires_replacement_approval( + chat_client_base: SupportsChatGetResponse, + streaming: bool, +) -> None: + """Middleware cannot silently execute repaired arguments under an approval for a different call.""" + executions: list[int] = [] + + class RepairArgumentsMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + assert isinstance(context.arguments, dict) + if "count_text" in context.arguments: + context.arguments = {"count": int(context.arguments["count_text"])} + await call_next() + + @tool(name="approved_count", approval_mode="always_require") + def approved_count(count: int) -> str: + executions.append(count) + return str(count) + + function_call = Content.from_function_call( + call_id="approved-repair-1", + name="approved_count", + arguments='{"count_text": "3"}', + ) + if streaming: + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + [ChatResponseUpdate(role="assistant", contents=[function_call])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")])], + ] + else: + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse(messages=Message(role="assistant", contents=[function_call])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + agent = Agent( + client=chat_client_base, + tools=[approved_count], + middleware=[RepairArgumentsMiddleware()], + ) + session = agent.create_session() + + async def run(value: str | Message): + if not streaming: + return await agent.run(value, session=session) + stream = agent.run(value, session=session, stream=True) + async for _ in stream: + pass + return await stream.get_final_response() + + first_response = await run("count") + first_request = next( + content + for message in first_response.messages + for content in message.contents + if content.type == "function_approval_request" + ) + replacement_response = await run( + Message(role="user", contents=[first_request.to_function_approval_response(approved=True)]) + ) + replacement_request = next( + content + for message in replacement_response.messages + for content in message.contents + if content.type == "function_approval_request" + ) + + assert executions == [] + assert replacement_request.id != first_request.id + assert replacement_request.additional_properties["_replacement_approval_request"] is True + assert replacement_request.function_call is not None + assert first_request.function_call is not None + assert replacement_request.function_call.id == first_request.function_call.id + assert replacement_request.function_call.parse_arguments() == {"count": 3} + + final_response = await run( + Message(role="user", contents=[replacement_request.to_function_approval_response(approved=True)]) + ) + + assert executions == [3] + assert final_response.text == "done" + + async def test_hosted_tool_approval_response(chat_client_base: SupportsChatGetResponse): """Test handling of approval responses for hosted tools (tools not in tool_map).""" diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index d8b500fde98..a69f53abc88 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -23,7 +23,7 @@ Message, SessionContext, ) -from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination +from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareFailure, MiddlewareTermination from agent_framework._tools import ( FunctionTool, _auto_invoke_function, @@ -5725,6 +5725,168 @@ async def execute(current: FunctionInvocationContext) -> list[Content]: 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}]"}, + ) + + result = await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration({"include_detailed_errors": include_detailed_errors}), + tool_map={strict_tool.name: strict_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) + + assert result.type == "function_result" + assert "secret payload" not in str(result.result) + assert "secret payload" not in str(result.exception) + assert result.exception == "Invalid arguments for 'strict_sink'. Invalid field(s): value." + + async def test_argument_mutation_after_security_middleware_fails_closed(self) -> None: + """A later middleware cannot change arguments after security policy has inspected them.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + sink = self._sink(accepts_untrusted=True) + + class LateRepairMiddleware(FunctionMiddleware): + async def process(self, context, call_next): + context.arguments = cast(Any, ["invalid", "post-policy", "arguments"]) + await call_next() + + function_call = Content.from_function_call( + call_id="late-repair", + name=sink.name, + arguments={"value": "approved value"}, + ) + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={sink.name: sink}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy, LateRepairMiddleware()), + ) + + async def test_argument_mutation_after_security_short_circuit_fails_closed(self) -> None: + """Post-policy mutation cannot evade the guard by short-circuiting execution.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + sink = self._sink(accepts_untrusted=True) + + class LateShortCircuitMiddleware(FunctionMiddleware): + async def process(self, context, call_next): + del call_next + context.arguments = {"value": "changed after policy"} + context.result = "short-circuited" + + function_call = Content.from_function_call( + call_id="late-short-circuit", + name=sink.name, + arguments={"value": "approved value"}, + ) + + with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): + await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={sink.name: sink}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy, LateShortCircuitMiddleware()), + ) + async def test_private_hidden_argument_is_blocked_from_public_sink(self) -> None: tracker = LabelTrackingFunctionMiddleware() policy = PolicyEnforcementFunctionMiddleware() diff --git a/python/samples/02-agents/tools/README.md b/python/samples/02-agents/tools/README.md index b24fe2223be..c4bfaf4bc4e 100644 --- a/python/samples/02-agents/tools/README.md +++ b/python/samples/02-agents/tools/README.md @@ -8,7 +8,7 @@ injection, and dynamic (progressive) tool exposure. | File | Demonstrates | |------|--------------| -| [`function_tool_with_explicit_schema.py`](function_tool_with_explicit_schema.py) | Defining a tool with an explicit JSON schema. | +| [`function_tool_with_explicit_schema.py`](function_tool_with_explicit_schema.py) | Choosing between Pydantic validation and a trusted, non-sensitive JSON schema declaration. | | [`function_tool_declaration_only.py`](function_tool_declaration_only.py) | A declaration-only tool (schema without a local implementation). | | [`function_tool_with_kwargs.py`](function_tool_with_kwargs.py) | Passing extra keyword arguments into a tool. | | [`function_tool_from_dict_with_dependency_injection.py`](function_tool_from_dict_with_dependency_injection.py) | Dependency injection into a tool defined from a dict. | @@ -16,6 +16,12 @@ injection, and dynamic (progressive) tool exposure. | [`tool_in_class.py`](tool_in_class.py) | Using a method on a class as a tool. | | [`agent_as_tool_with_session_propagation.py`](agent_as_tool_with_session_propagation.py) | Exposing an agent as a tool with session propagation. | +> [!WARNING] +> Caller-supplied JSON schema mappings receive only lightweight top-level runtime +> checks. Nested constraints and other JSON Schema keywords are not comprehensively +> enforced, so mappings must not be used as an authorization or security boundary. +> Use a Pydantic model for sensitive tools or whenever runtime validation matters. + ## Approvals & invocation control | File | Demonstrates | diff --git a/python/samples/02-agents/tools/function_tool_with_explicit_schema.py b/python/samples/02-agents/tools/function_tool_with_explicit_schema.py index 8090d4a2a02..9ef8b187afe 100644 --- a/python/samples/02-agents/tools/function_tool_with_explicit_schema.py +++ b/python/samples/02-agents/tools/function_tool_with_explicit_schema.py @@ -10,8 +10,14 @@ represent the desired schema. Two approaches are shown: -1. Using a Pydantic BaseModel subclass as the schema -2. Using a raw JSON schema dictionary as the schema +1. Using a Pydantic BaseModel subclass for recursive runtime validation +2. Using a raw JSON schema dictionary for a trusted, non-sensitive tool + +Raw schema dictionaries are passed through to the model and receive only +lightweight top-level checks at runtime. They do not comprehensively enforce +nested JSON Schema keywords and must not be used as an authorization or +security boundary. Prefer Pydantic for sensitive tools or whenever runtime +validation matters. """ import asyncio @@ -45,7 +51,8 @@ def get_weather(location: str, unit: str = "celsius") -> str: return f"The weather in {location} is 22 degrees {unit}." -# Approach 2: JSON schema dictionary as explicit schema +# Approach 2: JSON schema dictionary for a trusted, non-sensitive tool. +# This receives lightweight top-level checks only; use Pydantic for runtime enforcement. get_current_time_schema = { "type": "object", "properties": { From c15246c0ae711b33bddbc00cf8db9f24b62ce353 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 11 Sep 2026 11:54:58 +0200 Subject: [PATCH 2/4] Python: address argument validation review Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../specs/004-python-function-calling-loop.md | 4 +- .../packages/core/agent_framework/_tools.py | 44 ++++++---- .../core/test_function_invocation_logic.py | 78 ++++++++++++++++++ python/packages/core/tests/test_security.py | 81 ++++++++++++++++++- 4 files changed, 186 insertions(+), 21 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 426a78ed959..9f1b5da6d9b 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -597,8 +597,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Rejected execution | Rejection is a normal terminal result, not an exception to the caller. | `test_unapproved_tool_execution_raises_exception` | | Approved tool exception | Generic and detailed error modes preserve one result and one execution. | `test_approved_function_call_with_error_without_detailed_errors`, `test_approved_function_call_with_error_with_detailed_errors` | | Approved validation error | Validation failure returns one result without invoking the function body. | `test_approved_function_call_with_validation_error` | -| Pre-validation middleware repair | Function middleware observes raw parsed arguments, may repair them before final validation, and the body receives validated normalized values exactly once; short-circuiting skips validation and execution. Repair after security middleware fails closed, including invalid and short-circuited mutations. Validation errors after hidden-value resolution do not disclose the resolved value, while ordinary final Pydantic normalization remains allowed. | `test_function_middleware_repairs_raw_arguments_before_validation`, `test_function_middleware_can_short_circuit_before_argument_validation`, `test_invalid_arguments_produced_by_middleware_keep_argument_error_contract`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_argument_mutation_after_security_middleware_fails_closed`, `test_argument_mutation_after_security_short_circuit_fails_closed`, `test_hidden_argument_validation_error_does_not_disclose_resolved_value`, `test_hidden_argument_can_be_normalized_after_security_check` | -| Approved middleware repair | A changed approval-bound call executes zero times under the old grant, returns a persisted occurrence-bound replacement request in both response modes, and executes once only after the replacement is approved. Security expansion preserves approval-visible placeholders. | `test_approved_argument_repair_requires_replacement_approval`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_hidden_argument_resolution_does_not_require_reapproval` | +| Pre-validation middleware repair | Function middleware observes raw parsed arguments, may repair them before final validation, and the body receives validated normalized values exactly once; short-circuiting skips validation and execution. Repair after security middleware fails closed, including invalid and short-circuited mutations. Validation errors after hidden-value resolution do not disclose resolved values or mapping keys, including validator `TypeError` paths, while ordinary final Pydantic normalization remains allowed. | `test_function_middleware_repairs_raw_arguments_before_validation`, `test_function_middleware_can_short_circuit_before_argument_validation`, `test_invalid_arguments_produced_by_middleware_keep_argument_error_contract`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_argument_mutation_after_security_middleware_fails_closed`, `test_argument_mutation_after_security_short_circuit_fails_closed`, `test_hidden_argument_validation_error_does_not_disclose_resolved_value`, `test_hidden_mapping_key_is_not_disclosed_by_validation_error`, `test_hidden_value_is_not_disclosed_by_validator_type_error`, `test_hidden_argument_can_be_normalized_after_security_check` | +| Approved middleware repair | A changed approval-bound call executes zero times under the old grant, returns a persisted occurrence-bound replacement request in both response modes, and executes once only after the replacement is approved. The same replacement rule applies when middleware short-circuits instead of calling the tool. Security expansion preserves approval-visible placeholders. | `test_approved_argument_repair_requires_replacement_approval`, `test_approved_argument_repair_short_circuit_requires_replacement_approval`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_hidden_argument_resolution_does_not_require_reapproval` | | Approved success | Successful approved execution returns one result. | `test_approved_function_call_successful_execution` | | Consecutive error cap | Error threshold stops repeated failures, submits collected results, and makes only the required final no-tool model call. | `test_function_invocation_config_max_consecutive_errors`, `test_streaming_function_invocation_config_max_consecutive_errors`, `test_approval_resume_error_limit_forces_final_no_tool_response` | | Unknown call handling | Configured false returns an error result; configured true raises. | `test_function_invocation_config_terminate_on_unknown_calls_false`, `test_function_invocation_config_terminate_on_unknown_calls_true`, streaming equivalents | diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index c1299bab358..2540a21db3b 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -692,19 +692,15 @@ def _prepare_arguments(self, arguments: BaseModel | Mapping[str, Any] | None) -> f"Expected mapping-like arguments for tool '{self.name}', got {type(arguments).__name__}" ) except ValidationError as exc: - locations = sorted({ - ".".join(str(part) for part in error["loc"]) - for error in exc.errors(include_input=False, include_url=False) - }) - redacted_message = f"Invalid arguments for '{self.name}'." - if locations: - redacted_message = f"{redacted_message} Invalid field(s): {', '.join(locations)}." raise _FunctionArgumentValidationError( f"Invalid arguments for '{self.name}': {exc}", - redacted_message=redacted_message, + redacted_message=f"Invalid arguments for '{self.name}'.", ) from exc except TypeError as exc: - raise _FunctionArgumentValidationError(str(exc)) from exc + raise _FunctionArgumentValidationError( + str(exc), + redacted_message=f"Invalid arguments for '{self.name}'.", + ) from exc try: return _validate_arguments_against_schema( @@ -1917,6 +1913,20 @@ async def final_function_handler(context_obj: Any) -> Any: tool_call_id=call_id, ) + def ensure_short_circuit_arguments_are_authorized() -> None: + current_arguments = tool._arguments_as_mapping( # pyright: ignore[reportPrivateUsage] + middleware_context.arguments + ) + tool._ensure_security_arguments_unchanged( # pyright: ignore[reportPrivateUsage] + middleware_context, current_arguments + ) + tool._ensure_approved_arguments_unchanged( # pyright: ignore[reportPrivateUsage] + middleware_context, + tool._approval_visible_arguments( # pyright: ignore[reportPrivateUsage] + middleware_context.arguments, middleware_context + ), + ) + from ._middleware import MiddlewareTermination # MiddlewareTermination bubbles up to signal loop termination @@ -1926,10 +1936,7 @@ async def final_function_handler(context_obj: Any) -> Any: final_handler=final_function_handler, ) if not final_handler_started: - tool._ensure_security_arguments_unchanged( # pyright: ignore[reportPrivateUsage] - middleware_context, - tool._arguments_as_mapping(middleware_context.arguments), # pyright: ignore[reportPrivateUsage] - ) + ensure_short_circuit_arguments_are_authorized() # Pass through function_approval_request directly (e.g., from security middleware) if isinstance(function_result, Content) and function_result.type == "function_approval_request": @@ -1943,10 +1950,13 @@ async def final_function_handler(context_obj: Any) -> Any: ) except MiddlewareTermination as term_exc: if not final_handler_started: - tool._ensure_security_arguments_unchanged( # pyright: ignore[reportPrivateUsage] - middleware_context, - tool._arguments_as_mapping(middleware_context.arguments), # pyright: ignore[reportPrivateUsage] - ) + try: + ensure_short_circuit_arguments_are_authorized() + except _FunctionArgumentsChangedAfterApproval as exc: + raise MiddlewareTermination( + "Function arguments changed after approval.", + result=_replacement_approval_request(function_call_content, exc.arguments), + ) from exc # Re-raise to signal loop termination, but first capture any result set by middleware if middleware_context.result is not None: # Pass through function_approval_request directly (e.g., from security policy middleware) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 80a50de2d6a..7691fdb80ff 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3401,6 +3401,84 @@ async def run(value: str | Message): assert final_response.text == "done" +async def test_approved_argument_repair_short_circuit_requires_replacement_approval( + chat_client_base: SupportsChatGetResponse, +) -> None: + """Short-circuiting middleware cannot return under stale argument approval.""" + executions = 0 + + class RepairAndShortCircuitMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + del call_next + assert isinstance(context.arguments, dict) + if "count_text" in context.arguments: + context.arguments = {"count": int(context.arguments["count_text"])} + context.result = "handled by middleware" + + @tool(name="approved_short_circuit", approval_mode="always_require") + def approved_short_circuit(count: int) -> str: + nonlocal executions + executions += 1 + return str(count) + + function_call = Content.from_function_call( + call_id="approved-short-circuit-1", + name="approved_short_circuit", + arguments='{"count_text": "3"}', + ) + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse(messages=Message(role="assistant", contents=[function_call])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + agent = Agent( + client=chat_client_base, + tools=[approved_short_circuit], + middleware=[RepairAndShortCircuitMiddleware()], + ) + session = agent.create_session() + + first_response = await agent.run("count", session=session) + first_request = next( + content + for message in first_response.messages + for content in message.contents + if content.type == "function_approval_request" + ) + replacement_response = await agent.run( + Message(role="user", contents=[first_request.to_function_approval_response(approved=True)]), + session=session, + ) + replacement_request = next( + content + for message in replacement_response.messages + for content in message.contents + if content.type == "function_approval_request" + ) + + assert executions == 0 + assert replacement_request.function_call is not None + assert replacement_request.function_call.parse_arguments() == {"count": 3} + + final_response = await agent.run( + Message(role="user", contents=[replacement_request.to_function_approval_response(approved=True)]), + session=session, + ) + + assert executions == 0 + assert final_response.text == "done" + result = next( + content + for message in final_response.messages + for content in message.contents + if content.type == "function_result" + ) + assert result.result == "handled by middleware" + + async def test_hosted_tool_approval_response(chat_client_base: SupportsChatGetResponse): """Test handling of approval responses for hosted tools (tools not in tool_map).""" diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index a69f53abc88..2c9bf4b0add 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -11,7 +11,7 @@ from unittest.mock import AsyncMock import pytest -from pydantic import BaseModel +from pydantic import BaseModel, field_validator from agent_framework import ( Agent, @@ -5834,7 +5834,84 @@ class StrictArgs(BaseModel): assert result.type == "function_result" assert "secret payload" not in str(result.result) assert "secret payload" not in str(result.exception) - assert result.exception == "Invalid arguments for 'strict_sink'. Invalid field(s): value." + assert result.exception == "Invalid arguments for 'strict_sink'." + + async def test_hidden_mapping_key_is_not_disclosed_by_validation_error(self) -> None: + """Pydantic error locations cannot expose keys from resolved hidden mappings.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + {"secret-key": 1}, + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + + class StrictArgs(BaseModel): + value: dict[int, int] + + strict_tool = FunctionTool( + func=lambda value: str(value), + name="strict_mapping_sink", + input_model=StrictArgs, + additional_properties={"accepts_untrusted": True}, + ) + function_call = Content.from_function_call( + call_id="hidden-mapping-key-error", + name=strict_tool.name, + arguments={"value": f"[{variable_id}]"}, + ) + + result = await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration({"include_detailed_errors": True}), + tool_map={strict_tool.name: strict_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) + + assert result.type == "function_result" + assert "secret-key" not in str(result.result) + assert "secret-key" not in str(result.exception) + assert result.exception == "Invalid arguments for 'strict_mapping_sink'." + + async def test_hidden_value_is_not_disclosed_by_validator_type_error(self) -> None: + """Direct validator TypeErrors use the generic security redaction.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "secret payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + + class StrictArgs(BaseModel): + value: str + + @field_validator("value") + @classmethod + def reject_value(cls, value: str) -> str: + raise TypeError(f"rejected: {value}") + + strict_tool = FunctionTool( + func=lambda value: value, + name="validator_type_error_sink", + input_model=StrictArgs, + additional_properties={"accepts_untrusted": True}, + ) + function_call = Content.from_function_call( + call_id="hidden-validator-type-error", + name=strict_tool.name, + arguments={"value": f"[{variable_id}]"}, + ) + + result = await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration({"include_detailed_errors": True}), + tool_map={strict_tool.name: strict_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) + + assert result.type == "function_result" + assert "secret payload" not in str(result.result) + assert "secret payload" not in str(result.exception) + assert result.exception == "Invalid arguments for 'validator_type_error_sink'." async def test_argument_mutation_after_security_middleware_fails_closed(self) -> None: """A later middleware cannot change arguments after security policy has inspected them.""" From 5d0560975733ddc16404cb532df0deda3e8244ac Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 11 Sep 2026 13:43:37 +0200 Subject: [PATCH 3/4] Python: preserve validated middleware arguments Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../specs/004-python-function-calling-loop.md | 20 ++-- .../core/agent_framework/_middleware.py | 41 ++++--- .../packages/core/agent_framework/_tools.py | 84 +++++++++++++- .../packages/core/agent_framework/security.py | 7 ++ .../core/test_function_invocation_logic.py | 104 +++++++++++++++++- python/packages/core/tests/test_security.py | 59 +++++++++- .../02-agents/devui/agent_weather/agent.py | 8 +- 7 files changed, 288 insertions(+), 35 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 9f1b5da6d9b..14a3f593f49 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -377,13 +377,15 @@ that manually replay messages own the equivalent rule: do not resend an approval not invent one. - A completed function call/result pair is inert on later turns. - Informational-only and declaration-only calls are not executed as local tools. -- For automatic local execution, provider arguments are JSON-parsed before function middleware but are not schema - validated yet. Middleware may inspect or repair that raw mapping before calling `call_next()`. The innermost handler - performs final validation immediately before the tool body and writes the validated, normalized mapping back to - `FunctionInvocationContext.arguments`. Middleware that short-circuits without `call_next()` also skips final - validation and tool execution. -- Argument-repair middleware must precede security or policy middleware so enforcement observes the effective - invocation. Changing arguments after security middleware has processed them fails closed with `MiddlewareFailure`. +- For automatic local execution, provider arguments are JSON-parsed before function middleware. Schema-compatible + arguments retain the existing normalized mapping contract. A provisional validation failure is not terminal: + middleware receives the raw mapping and may repair it before calling `call_next()`. The innermost handler validates + changed or previously invalid arguments immediately before the tool body and writes the normalized mapping back to + `FunctionInvocationContext.arguments`; unchanged provisionally normalized arguments are reused without running + validators twice. Middleware that short-circuits without `call_next()` skips final validation and tool execution. +- Argument-repair middleware must precede security or policy middleware so enforcement observes the effective, + normalized invocation. Built-in security middleware normalizes hidden-value expansions before policy inspection. + Changing arguments after security middleware has processed them fails closed with `MiddlewareFailure`. - Argument-validation failures after middleware retain the established `Argument parsing failed` result contract. Exceptions raised by middleware or the tool body retain the separate `Function failed` contract. @@ -597,8 +599,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Rejected execution | Rejection is a normal terminal result, not an exception to the caller. | `test_unapproved_tool_execution_raises_exception` | | Approved tool exception | Generic and detailed error modes preserve one result and one execution. | `test_approved_function_call_with_error_without_detailed_errors`, `test_approved_function_call_with_error_with_detailed_errors` | | Approved validation error | Validation failure returns one result without invoking the function body. | `test_approved_function_call_with_validation_error` | -| Pre-validation middleware repair | Function middleware observes raw parsed arguments, may repair them before final validation, and the body receives validated normalized values exactly once; short-circuiting skips validation and execution. Repair after security middleware fails closed, including invalid and short-circuited mutations. Validation errors after hidden-value resolution do not disclose resolved values or mapping keys, including validator `TypeError` paths, while ordinary final Pydantic normalization remains allowed. | `test_function_middleware_repairs_raw_arguments_before_validation`, `test_function_middleware_can_short_circuit_before_argument_validation`, `test_invalid_arguments_produced_by_middleware_keep_argument_error_contract`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_argument_mutation_after_security_middleware_fails_closed`, `test_argument_mutation_after_security_short_circuit_fails_closed`, `test_hidden_argument_validation_error_does_not_disclose_resolved_value`, `test_hidden_mapping_key_is_not_disclosed_by_validation_error`, `test_hidden_value_is_not_disclosed_by_validator_type_error`, `test_hidden_argument_can_be_normalized_after_security_check` | -| Approved middleware repair | A changed approval-bound call executes zero times under the old grant, returns a persisted occurrence-bound replacement request in both response modes, and executes once only after the replacement is approved. The same replacement rule applies when middleware short-circuits instead of calling the tool. Security expansion preserves approval-visible placeholders. | `test_approved_argument_repair_requires_replacement_approval`, `test_approved_argument_repair_short_circuit_requires_replacement_approval`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_hidden_argument_resolution_does_not_require_reapproval` | +| Pre-validation middleware repair | Schema-compatible calls retain normalized middleware arguments without duplicate validation. When provisional validation fails, function middleware observes raw parsed arguments, may repair them before final validation, and the body receives normalized values; short-circuiting skips final validation and execution. Repair after security middleware fails closed using recursive type-aware comparison, including invalid and short-circuited mutations. Security inspects exact normalized values, and validation errors after hidden-value resolution do not disclose resolved values or mapping keys, including validator `TypeError` paths. | `test_function_middleware_keeps_normalized_arguments_for_valid_calls`, `test_function_middleware_repairs_raw_arguments_before_validation`, `test_function_middleware_can_short_circuit_before_argument_validation`, `test_invalid_arguments_produced_by_middleware_keep_argument_error_contract`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_argument_mutation_after_security_middleware_fails_closed`, `test_argument_mutation_after_security_short_circuit_fails_closed`, `test_security_policy_observes_custom_validator_transform_once`, `test_hidden_argument_validation_error_does_not_disclose_resolved_value`, `test_hidden_mapping_key_is_not_disclosed_by_validation_error`, `test_hidden_value_is_not_disclosed_by_validator_type_error`, `test_hidden_argument_can_be_normalized_after_security_check` | +| Approved middleware repair | A changed approval-bound call executes zero times under the old grant, returns a persisted occurrence-bound replacement request in both response modes, and executes once only after the replacement is approved. Recursive type-aware comparison treats booleans and numbers as distinct. The same replacement rule applies when middleware short-circuits instead of calling the tool. Security expansion preserves approval-visible placeholders. | `test_approved_argument_repair_requires_replacement_approval`, `test_approved_argument_repair_short_circuit_requires_replacement_approval`, `test_approval_snapshot_distinguishes_boolean_from_integer`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_hidden_argument_resolution_does_not_require_reapproval` | | Approved success | Successful approved execution returns one result. | `test_approved_function_call_successful_execution` | | Consecutive error cap | Error threshold stops repeated failures, submits collected results, and makes only the required final no-tool model call. | `test_function_invocation_config_max_consecutive_errors`, `test_streaming_function_invocation_config_max_consecutive_errors`, `test_approval_resume_error_limit_forces_final_no_tool_response` | | Unknown call handling | Configured false returns an error result; configured true raises. | `test_function_invocation_config_terminate_on_unknown_calls_false`, `test_function_invocation_config_terminate_on_unknown_calls_true`, streaming equivalents | diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 1423ec83b76..bc1f4610a75 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -371,12 +371,14 @@ class FunctionInvocationContext: Attributes: function: The function being invoked. arguments: The function arguments. In the automatic function-calling loop, - middleware initially receives the raw JSON-parsed mapping from the - provider and may repair it before calling ``call_next()``. The innermost - handler validates and normalizes the current value immediately before + schema-compatible provider arguments retain the existing normalized + mapping contract. If provisional normalization rejects provider + arguments, middleware instead receives the raw JSON-parsed mapping and + may repair it before calling ``call_next()``. The innermost handler + validates changed or previously invalid arguments immediately before execution, then stores the normalized mapping back on this attribute. - Middleware that short-circuits without calling ``call_next()`` also skips - final schema validation and function execution. + Middleware that short-circuits without calling ``call_next()`` skips + final validation and function execution. session: The agent session for this invocation, if any. metadata: Metadata dictionary for sharing data between function middleware. result: Function execution result. This attribute carries no guaranteed type. @@ -444,9 +446,9 @@ def __init__( Args: function: The function being invoked. - arguments: The function arguments. Automatic invocation supplies the raw - JSON-parsed mapping; final validation occurs after middleware, immediately - before function execution. + arguments: The function arguments. Automatic invocation supplies a normalized + mapping when provisional validation succeeds, otherwise the raw JSON-parsed + mapping so middleware can repair it before final validation. session: The agent session for this invocation, if any. metadata: Metadata dictionary for sharing data between function middleware. result: Function execution result. Observed and overridden values do not @@ -722,11 +724,12 @@ class FunctionMiddleware(ABC): """Abstract base class for function middleware that can intercept function invocations. Function middleware allows you to intercept and modify function/tool invocations before - and after execution. On entry, automatic function invocation exposes the raw JSON-parsed - arguments so middleware can repair provider-specific deviations before calling - ``call_next()``. The innermost handler performs final validation immediately before - execution and updates ``context.arguments`` with the normalized values. You can also - cache results, log invocations, or override function execution. + and after execution. On entry, schema-compatible calls retain normalized arguments. + When provisional normalization rejects provider arguments, middleware receives the raw + JSON-parsed mapping so it can repair provider-specific deviations before calling + ``call_next()``. The innermost handler validates changed or previously invalid arguments + immediately before execution and updates ``context.arguments`` with normalized values. + You can also cache results, log invocations, or override function execution. Argument-repair middleware must run before security or policy middleware so those layers inspect the effective invocation. Changing arguments after security middleware @@ -783,11 +786,13 @@ async def process( Args: context: Function invocation context containing function, arguments, and metadata. - Before ``call_next()``, automatic invocation exposes raw JSON-parsed - arguments that middleware may inspect or replace. After ``call_next()`` - reaches the function, arguments contain their validated, normalized - values. MiddlewareTypes can set context.result to override execution, - or observe the actual execution result after calling call_next(). + Before ``call_next()``, automatic invocation exposes normalized + arguments for schema-compatible calls and raw JSON-parsed arguments + when provisional validation failed. Middleware may inspect or replace + either mapping. After ``call_next()`` reaches the function, arguments + contain their validated, normalized values. MiddlewareTypes can set + context.result to override execution, or observe the actual execution + result after calling call_next(). call_next: Function to call the next middleware or final function execution. Does not return anything - all data flows through the context. diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 2540a21db3b..392a601f8a3 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -123,6 +123,8 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None) _FUNCTION_RESULT_PAYLOAD_BUDGET_STATE_KEY: Final[str] = "_function_result_payload_budget" _APPROVED_ARGUMENTS_CONTEXT_KEY: Final[str] = "_approved_function_arguments" _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" _FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT: Final[str] = ( "Function invocation limit reached before a final answer could be produced." ) @@ -145,6 +147,41 @@ def __init__(self, arguments: Mapping[str, Any]) -> None: self.arguments = copy.deepcopy(dict(arguments)) +def _type_aware_equal(left: Any, right: Any) -> bool: + """Compare nested argument values without collapsing distinct JSON types.""" + if type(left) is not type(right): + return False + if isinstance(left, dict): + left_dict = cast(dict[Any, Any], left) + right_dict = cast(dict[Any, Any], right) + right_items: list[tuple[Any, Any]] = list(right_dict.items()) + if len(left_dict) != len(right_items): + return False + for left_key, left_value in left_dict.items(): + match_index = next( + ( + index + for index, (right_key, _) in enumerate(right_items) + if type(left_key) is type(right_key) and left_key == right_key + ), + None, + ) + if match_index is None: + return False + _, right_value = right_items.pop(match_index) + if not _type_aware_equal(left_value, right_value): + return False + return True + if isinstance(left, list | tuple): + left_sequence = cast(list[Any] | tuple[Any, ...], left) + right_sequence = cast(list[Any] | tuple[Any, ...], right) + return len(left_sequence) == len(right_sequence) and all( + _type_aware_equal(left_item, right_item) + for left_item, right_item in zip(left_sequence, right_sequence, strict=True) + ) + return bool(left == right) + + ApprovalMode: TypeAlias = Literal["always_require", "never_require"] ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]") ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) @@ -734,6 +771,26 @@ def _approval_visible_arguments( return cls._arguments_as_mapping(context.metadata["original_arguments_for_messages"]) return cls._arguments_as_mapping(arguments) + def _prepare_context_arguments( + self, + context: FunctionInvocationContext, + arguments: BaseModel | Mapping[str, Any] | None, + ) -> dict[str, Any]: + """Prepare current context arguments once, reusing an unchanged prepared snapshot.""" + current_arguments = self._arguments_as_mapping(arguments) + prepared_arguments = context.metadata.get(_PREPARED_ARGUMENTS_CONTEXT_KEY) + if ( + current_arguments is not None + and isinstance(prepared_arguments, Mapping) + and _type_aware_equal(current_arguments, dict(cast(Mapping[str, Any], prepared_arguments))) + ): + return current_arguments + + validated_arguments = self._prepare_arguments(arguments) + context.arguments = validated_arguments + context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(validated_arguments) + return validated_arguments + @staticmethod def _ensure_security_arguments_unchanged( context: FunctionInvocationContext | None, @@ -744,7 +801,8 @@ def _ensure_security_arguments_unchanged( return security_arguments = context.metadata.get(_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY) if isinstance(security_arguments, Mapping) and ( - current_arguments is None or dict(cast(Mapping[str, Any], security_arguments)) != dict(current_arguments) + current_arguments is None + or not _type_aware_equal(dict(cast(Mapping[str, Any], security_arguments)), dict(current_arguments)) ): from ._middleware import MiddlewareFailure @@ -766,7 +824,10 @@ def _ensure_approved_arguments_unchanged( approved_arguments = context.metadata.get(_APPROVED_ARGUMENTS_CONTEXT_KEY) if not isinstance(approved_arguments, Mapping): return - if dict(cast(Mapping[str, Any], approved_arguments)) != dict(approval_visible_arguments): + if not _type_aware_equal( + dict(cast(Mapping[str, Any], approved_arguments)), + dict(approval_visible_arguments), + ): raise _FunctionArgumentsChangedAfterApproval(approval_visible_arguments) @overload @@ -863,7 +924,11 @@ async def invoke( current_arguments = self._arguments_as_mapping(arguments) approval_visible_arguments = self._approval_visible_arguments(arguments, context) self._ensure_security_arguments_unchanged(context, current_arguments) - validated_arguments = self._prepare_arguments(arguments) + validated_arguments = ( + self._prepare_context_arguments(context, arguments) + if context is not None + else self._prepare_arguments(arguments) + ) effective_context = context if effective_context is None and self._context_parameter_name is not None: @@ -876,6 +941,7 @@ async def invoke( effective_context.function = self effective_context.arguments = validated_arguments effective_context.kwargs = dict(runtime_kwargs) + effective_context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(validated_arguments) self._ensure_approved_arguments_unchanged(effective_context, approval_visible_arguments) @@ -1877,6 +1943,15 @@ async def _auto_invoke_function( except Exception as exc: return _function_execution_error_result(function_call_content, tool.name, exc, config, direct_context) # Execute through middleware pipeline if available + arguments_prepared = False + try: + args = tool._prepare_arguments(args) # pyright: ignore[reportPrivateUsage] + arguments_prepared = True + except _FunctionArgumentValidationError: + # Invalid provider arguments are intentionally exposed to middleware so + # it has a supported opportunity to repair them before final validation. + pass + middleware_context = FunctionInvocationContext( function=tool, arguments=args, @@ -1886,6 +1961,9 @@ async def _auto_invoke_function( ) if host_payload_budget is not None: middleware_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget + middleware_context.metadata[_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY] = True + if arguments_prepared: + middleware_context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(args) call_id = function_call_content.call_id if call_id is None: diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index c4acbf33d3f..1a36257db9c 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -38,6 +38,7 @@ from ._sessions import AgentSession, ContextProvider from ._tools import ( _APPROVAL_REQUEST_ID_KEY, # pyright: ignore[reportPrivateUsage] + _AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] FunctionTool, tool, @@ -1555,6 +1556,12 @@ async def process( # 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] = deepcopy(context.arguments) + if context.metadata.get(_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY) is True: + context.function._prepare_context_arguments( # pyright: ignore[reportPrivateUsage] + context, + context.arguments, + ) + context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = deepcopy(context.arguments) argument_labels = [*input_labels, *resolved_labels] argument_label = combine_labels(*argument_labels) if argument_labels else ContentLabel() diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 7691fdb80ff..1115252db87 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -8,6 +8,7 @@ from typing import Any, Literal import pytest +from pydantic import BaseModel, field_validator from agent_framework import ( Agent, @@ -32,7 +33,12 @@ annotate_message_groups, included_token_count, ) -from agent_framework._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination +from agent_framework._middleware import ( + FunctionInvocationContext, + FunctionMiddleware, + FunctionMiddlewarePipeline, + MiddlewareTermination, +) _EXPECTED_FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT = ( "Function invocation limit reached before a final answer could be produced." @@ -3172,6 +3178,59 @@ def count_items(count: int) -> str: assert executed_arguments == [3] +async def test_function_middleware_keeps_normalized_arguments_for_valid_calls( + chat_client_base: SupportsChatGetResponse, +) -> None: + """Valid calls preserve the existing normalized middleware argument contract.""" + validation_count = 0 + observed_arguments: list[dict[str, Any]] = [] + + class CountArgs(BaseModel): + count: int + + @field_validator("count") + @classmethod + def track_validation(cls, value: int) -> int: + nonlocal validation_count + validation_count += 1 + return value + + class ObserveArgumentsMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + assert isinstance(context.arguments, dict) + observed_arguments.append(dict(context.arguments)) + await call_next() + + @tool(name="count_items", schema=CountArgs, approval_mode="never_require") + def count_items(count: int) -> str: + return str(count) + + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="valid-1", name="count_items", arguments='{"count": "3"}') + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + await chat_client_base.get_response( + [Message(role="user", contents=["count"])], + options={"tools": [count_items]}, + client_kwargs={"middleware": [ObserveArgumentsMiddleware()]}, + ) + + assert observed_arguments == [{"count": 3}] + assert validation_count == 1 + + async def test_function_middleware_can_short_circuit_before_argument_validation( chat_client_base: SupportsChatGetResponse, ) -> None: @@ -3479,6 +3538,49 @@ def approved_short_circuit(count: int) -> str: assert result.result == "handled by middleware" +async def test_approval_snapshot_distinguishes_boolean_from_integer() -> None: + """A boolean-to-integer mutation requires replacement approval.""" + from agent_framework._tools import _auto_invoke_function, normalize_function_invocation_configuration + + class ChangeTypeMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.arguments = {"value": 1} + await call_next() + + @tool(name="typed_value", approval_mode="always_require") + def typed_value(value: Any) -> str: + return f"{type(value).__name__}:{value}" + + function_call = Content.from_function_call( + call_id="typed-approval", + id="typed-approval-occurrence", + name=typed_value.name, + arguments={"value": True}, + ) + approval_response = Content.from_function_approval_request( + id="typed-approval-occurrence", + function_call=function_call, + ).to_function_approval_response(approved=True) + + with pytest.raises(MiddlewareTermination) as exc_info: + await _auto_invoke_function( + approval_response, + config=normalize_function_invocation_configuration(None), + tool_map={typed_value.name: typed_value}, + middleware_pipeline=FunctionMiddlewarePipeline(ChangeTypeMiddleware()), + ) + + replacement = exc_info.value.result + assert isinstance(replacement, Content) + assert replacement.type == "function_approval_request" + assert replacement.function_call is not None + assert replacement.function_call.parse_arguments() == {"value": 1} + + async def test_hosted_tool_approval_response(chat_client_base: SupportsChatGetResponse): """Test handling of approval responses for hosted tools (tools not in tool_map).""" diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 2c9bf4b0add..16bdfad8618 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -5921,13 +5921,13 @@ async def test_argument_mutation_after_security_middleware_fails_closed(self) -> class LateRepairMiddleware(FunctionMiddleware): async def process(self, context, call_next): - context.arguments = cast(Any, ["invalid", "post-policy", "arguments"]) + context.arguments = {"value": 1} await call_next() function_call = Content.from_function_call( call_id="late-repair", name=sink.name, - arguments={"value": "approved value"}, + arguments={"value": True}, ) with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): @@ -5938,6 +5938,61 @@ async def process(self, context, call_next): middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy, LateRepairMiddleware()), ) + async def test_security_policy_observes_custom_validator_transform_once(self) -> None: + """Security middleware inspects the exact normalized value delivered to the tool.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + validation_count = 0 + observed: list[str] = [] + received: list[str] = [] + + class TransformArgs(BaseModel): + value: str + + @field_validator("value") + @classmethod + def transform_value(cls, value: str) -> str: + nonlocal validation_count + validation_count += 1 + return "dangerous-operation" if value == "safe" else value + + class ObserveAfterTrackingMiddleware(FunctionMiddleware): + async def process(self, context, call_next): + observed.append(cast(dict[str, str], context.arguments)["value"]) + await call_next() + + def transformed_sink(value: str) -> str: + received.append(value) + return value + + transformed_tool = FunctionTool( + func=transformed_sink, + name="transformed_sink", + input_model=TransformArgs, + additional_properties={"accepts_untrusted": True}, + ) + function_call = Content.from_function_call( + call_id="validator-transform", + name=transformed_tool.name, + arguments={"value": "safe"}, + ) + + result = await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={transformed_tool.name: transformed_tool}, + middleware_pipeline=FunctionMiddlewarePipeline( + tracker, + ObserveAfterTrackingMiddleware(), + policy, + ), + ) + + assert result.type == "function_result" + assert observed == ["dangerous-operation"] + assert received == ["dangerous-operation"] + assert validation_count == 1 + async def test_argument_mutation_after_security_short_circuit_fails_closed(self) -> None: """Post-policy mutation cannot evade the guard by short-circuiting execution.""" tracker = LabelTrackingFunctionMiddleware() diff --git a/python/samples/02-agents/devui/agent_weather/agent.py b/python/samples/02-agents/devui/agent_weather/agent.py index 201053ab923..b2a316ae6bd 100644 --- a/python/samples/02-agents/devui/agent_weather/agent.py +++ b/python/samples/02-agents/devui/agent_weather/agent.py @@ -3,7 +3,7 @@ import logging import os -from collections.abc import AsyncIterable, Awaitable, Callable +from collections.abc import AsyncIterable, Awaitable, Callable, Mapping from typing import Annotated from agent_framework import ( @@ -93,7 +93,11 @@ async def atlantis_location_filter_middleware( ) -> None: """Function middleware that blocks weather requests for Atlantis.""" # Check if location parameter is "atlantis" - location = getattr(context.arguments, "location", None) + location = ( + context.arguments.get("location") + if isinstance(context.arguments, Mapping) + else getattr(context.arguments, "location", None) + ) if location and location.lower() == "atlantis": context.result = ( "Blocked! Hold up right there!! Tell the user that " From 535e269af452d164609fdf60caba78e805cd2145 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 11 Sep 2026 16:19:15 +0200 Subject: [PATCH 4/4] Python: harden function argument snapshots Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../specs/004-python-function-calling-loop.md | 4 +- .../packages/core/agent_framework/_tools.py | 103 ++++++----- .../packages/core/agent_framework/security.py | 7 +- .../core/test_function_invocation_logic.py | 172 +++++++++++++++++- python/packages/core/tests/test_security.py | 53 +++++- 5 files changed, 274 insertions(+), 65 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 14a3f593f49..6b7da95f51c 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -599,8 +599,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Rejected execution | Rejection is a normal terminal result, not an exception to the caller. | `test_unapproved_tool_execution_raises_exception` | | Approved tool exception | Generic and detailed error modes preserve one result and one execution. | `test_approved_function_call_with_error_without_detailed_errors`, `test_approved_function_call_with_error_with_detailed_errors` | | Approved validation error | Validation failure returns one result without invoking the function body. | `test_approved_function_call_with_validation_error` | -| Pre-validation middleware repair | Schema-compatible calls retain normalized middleware arguments without duplicate validation. When provisional validation fails, function middleware observes raw parsed arguments, may repair them before final validation, and the body receives normalized values; short-circuiting skips final validation and execution. Repair after security middleware fails closed using recursive type-aware comparison, including invalid and short-circuited mutations. Security inspects exact normalized values, and validation errors after hidden-value resolution do not disclose resolved values or mapping keys, including validator `TypeError` paths. | `test_function_middleware_keeps_normalized_arguments_for_valid_calls`, `test_function_middleware_repairs_raw_arguments_before_validation`, `test_function_middleware_can_short_circuit_before_argument_validation`, `test_invalid_arguments_produced_by_middleware_keep_argument_error_contract`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_argument_mutation_after_security_middleware_fails_closed`, `test_argument_mutation_after_security_short_circuit_fails_closed`, `test_security_policy_observes_custom_validator_transform_once`, `test_hidden_argument_validation_error_does_not_disclose_resolved_value`, `test_hidden_mapping_key_is_not_disclosed_by_validation_error`, `test_hidden_value_is_not_disclosed_by_validator_type_error`, `test_hidden_argument_can_be_normalized_after_security_check` | -| Approved middleware repair | A changed approval-bound call executes zero times under the old grant, returns a persisted occurrence-bound replacement request in both response modes, and executes once only after the replacement is approved. Recursive type-aware comparison treats booleans and numbers as distinct. The same replacement rule applies when middleware short-circuits instead of calling the tool. Security expansion preserves approval-visible placeholders. | `test_approved_argument_repair_requires_replacement_approval`, `test_approved_argument_repair_short_circuit_requires_replacement_approval`, `test_approval_snapshot_distinguishes_boolean_from_integer`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_hidden_argument_resolution_does_not_require_reapproval` | +| Pre-validation middleware repair | Schema-compatible calls retain normalized middleware arguments without duplicate validation. Prepared-value reuse does not require validator outputs to be copyable, and unchanged NaNs remain stable. When provisional validation fails, function middleware observes raw parsed arguments, may repair them before final validation, and the body receives normalized values; short-circuiting skips final validation and execution. Repair after security middleware fails closed using recursive type-aware, float-bit-exact comparison, including invalid and short-circuited mutations. Security inspects exact normalized values, and validation errors after hidden-value resolution do not disclose resolved values or mapping keys, including validator `TypeError` paths. | `test_function_middleware_keeps_normalized_arguments_for_valid_calls`, `test_prepared_arguments_support_noncopyable_validator_output`, `test_nan_prepared_and_approval_snapshots_are_stable`, `test_function_middleware_repairs_raw_arguments_before_validation`, `test_function_middleware_can_short_circuit_before_argument_validation`, `test_invalid_arguments_produced_by_middleware_keep_argument_error_contract`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_argument_mutation_after_security_middleware_fails_closed`, `test_security_snapshot_accepts_unchanged_nan`, `test_argument_mutation_after_security_short_circuit_fails_closed`, `test_security_policy_observes_custom_validator_transform_once`, `test_hidden_argument_validation_error_does_not_disclose_resolved_value`, `test_hidden_mapping_key_is_not_disclosed_by_validation_error`, `test_hidden_value_is_not_disclosed_by_validator_type_error`, `test_hidden_argument_can_be_normalized_after_security_check` | +| Approved middleware repair | Approval binds to the normalized middleware-entry representation, so ordinary Pydantic coercion still completes in one approval round. A changed approval-bound call executes zero times under the old grant, returns a persisted occurrence-bound replacement request in both response modes, and executes once only after the replacement is approved. Recursive type-aware, float-bit-exact comparison treats booleans and numbers, and positive and negative zero, as distinct while keeping unchanged NaNs stable. The same replacement rule applies when middleware short-circuits instead of calling the tool. Security expansion preserves approval-visible placeholders. | `test_approved_coercing_arguments_execute_without_replacement`, `test_approved_argument_repair_requires_replacement_approval`, `test_approved_argument_repair_short_circuit_requires_replacement_approval`, `test_approval_snapshot_distinguishes_exact_values`, `packages/core/tests/test_security.py::TestVariableArgumentPolicy::test_hidden_argument_resolution_does_not_require_reapproval` | | Approved success | Successful approved execution returns one result. | `test_approved_function_call_successful_execution` | | Consecutive error cap | Error threshold stops repeated failures, submits collected results, and makes only the required final no-tool model call. | `test_function_invocation_config_max_consecutive_errors`, `test_streaming_function_invocation_config_max_consecutive_errors`, `test_approval_resume_error_limit_forces_final_no_tool_response` | | Unknown call handling | Configured false returns an error result; configured true raises. | `test_function_invocation_config_terminate_on_unknown_calls_false`, `test_function_invocation_config_terminate_on_unknown_calls_true`, streaming equivalents | diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index c30a4f9542a..24fb7fd5378 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -8,6 +8,7 @@ import inspect import json import logging +import struct import sys import typing import warnings @@ -144,42 +145,38 @@ class _FunctionArgumentsChangedAfterApproval(Exception): def __init__(self, arguments: Mapping[str, Any]) -> None: super().__init__("Function arguments changed after approval.") - self.arguments = copy.deepcopy(dict(arguments)) + self.arguments = dict(arguments) + + +def _argument_comparison_token(value: Any) -> Any: + """Build an immutable, type-aware token without copying argument objects.""" + if isinstance(value, BaseModel): + return _argument_comparison_token(value.model_dump(exclude_unset=True)) + if isinstance(value, dict): + return ( + "dict", + frozenset( + (_argument_comparison_token(key), _argument_comparison_token(item)) + for key, item in cast(dict[Any, Any], value).items() + ), + ) + if isinstance(value, list): + return ("list", tuple(_argument_comparison_token(item) for item in cast(list[Any], value))) + if isinstance(value, tuple): + return ("tuple", tuple(_argument_comparison_token(item) for item in cast(tuple[Any, ...], value))) + if isinstance(value, float): + return ("float", struct.pack("!d", value)) + if value is None or isinstance(value, bool | int | str | bytes): + return (type(value), value) + return ("object", type(value), id(value)) -def _type_aware_equal(left: Any, right: Any) -> bool: - """Compare nested argument values without collapsing distinct JSON types.""" - if type(left) is not type(right): - return False - if isinstance(left, dict): - left_dict = cast(dict[Any, Any], left) - right_dict = cast(dict[Any, Any], right) - right_items: list[tuple[Any, Any]] = list(right_dict.items()) - if len(left_dict) != len(right_items): - return False - for left_key, left_value in left_dict.items(): - match_index = next( - ( - index - for index, (right_key, _) in enumerate(right_items) - if type(left_key) is type(right_key) and left_key == right_key - ), - None, - ) - if match_index is None: - return False - _, right_value = right_items.pop(match_index) - if not _type_aware_equal(left_value, right_value): - return False - return True - if isinstance(left, list | tuple): - left_sequence = cast(list[Any] | tuple[Any, ...], left) - right_sequence = cast(list[Any] | tuple[Any, ...], right) - return len(left_sequence) == len(right_sequence) and all( - _type_aware_equal(left_item, right_item) - for left_item, right_item in zip(left_sequence, right_sequence, strict=True) - ) - return bool(left == right) +@dataclass(frozen=True) +class _PreparedArgumentsState: + """Exact prepared arguments and their immutable comparison token.""" + + arguments: dict[str, Any] + token: Any ApprovalMode: TypeAlias = Literal["always_require", "never_require"] @@ -784,17 +781,21 @@ def _prepare_context_arguments( ) -> dict[str, Any]: """Prepare current context arguments once, reusing an unchanged prepared snapshot.""" current_arguments = self._arguments_as_mapping(arguments) - prepared_arguments = context.metadata.get(_PREPARED_ARGUMENTS_CONTEXT_KEY) + prepared_state = context.metadata.get(_PREPARED_ARGUMENTS_CONTEXT_KEY) if ( current_arguments is not None - and isinstance(prepared_arguments, Mapping) - and _type_aware_equal(current_arguments, dict(cast(Mapping[str, Any], prepared_arguments))) + and isinstance(prepared_state, _PreparedArgumentsState) + and _argument_comparison_token(current_arguments) == prepared_state.token ): - return current_arguments + context.arguments = prepared_state.arguments + return prepared_state.arguments validated_arguments = self._prepare_arguments(arguments) context.arguments = validated_arguments - context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(validated_arguments) + context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = _PreparedArgumentsState( + arguments=validated_arguments, + token=_argument_comparison_token(validated_arguments), + ) return validated_arguments @staticmethod @@ -805,10 +806,9 @@ def _ensure_security_arguments_unchanged( """Fail closed when arguments change after security middleware processed them.""" if context is None: return - security_arguments = context.metadata.get(_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY) - if isinstance(security_arguments, Mapping) and ( - current_arguments is None - or not _type_aware_equal(dict(cast(Mapping[str, Any], security_arguments)), dict(current_arguments)) + security_token = context.metadata.get(_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY) + if security_token is not None and ( + current_arguments is None or _argument_comparison_token(dict(current_arguments)) != security_token ): from ._middleware import MiddlewareFailure @@ -827,13 +827,10 @@ def _ensure_approved_arguments_unchanged( return if approval_visible_arguments is None: return - approved_arguments = context.metadata.get(_APPROVED_ARGUMENTS_CONTEXT_KEY) - if not isinstance(approved_arguments, Mapping): + approved_token = context.metadata.get(_APPROVED_ARGUMENTS_CONTEXT_KEY) + if approved_token is None: return - if not _type_aware_equal( - dict(cast(Mapping[str, Any], approved_arguments)), - dict(approval_visible_arguments), - ): + if _argument_comparison_token(dict(approval_visible_arguments)) != approved_token: raise _FunctionArgumentsChangedAfterApproval(approval_visible_arguments) @overload @@ -947,7 +944,6 @@ async def invoke( effective_context.function = self effective_context.arguments = validated_arguments effective_context.kwargs = dict(runtime_kwargs) - effective_context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(validated_arguments) self._ensure_approved_arguments_unchanged(effective_context, approval_visible_arguments) @@ -1969,7 +1965,10 @@ async def _auto_invoke_function( middleware_context.metadata[_FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY] = host_payload_budget middleware_context.metadata[_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY] = True if arguments_prepared: - middleware_context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(args) + middleware_context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = _PreparedArgumentsState( + arguments=args, + token=_argument_comparison_token(args), + ) call_id = function_call_content.call_id if call_id is None: @@ -1984,7 +1983,7 @@ async def _auto_invoke_function( # this replay corresponds to a middleware-specific approval flow. if approval_response is not None: middleware_context.metadata["approval_response"] = approval_response - middleware_context.metadata[_APPROVED_ARGUMENTS_CONTEXT_KEY] = copy.deepcopy(parsed_args) + middleware_context.metadata[_APPROVED_ARGUMENTS_CONTEXT_KEY] = _argument_comparison_token(args) final_handler_started = False diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 1a36257db9c..f4cda8e9b34 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -41,6 +41,7 @@ _AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY, # pyright: ignore[reportPrivateUsage] FunctionTool, + _argument_comparison_token, # pyright: ignore[reportPrivateUsage] tool, ) from ._types import Content, Message @@ -1555,13 +1556,15 @@ async def process( # 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] = deepcopy(context.arguments) + context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = _argument_comparison_token(context.arguments) if context.metadata.get(_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY) is True: context.function._prepare_context_arguments( # pyright: ignore[reportPrivateUsage] context, context.arguments, ) - context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = deepcopy(context.arguments) + context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = _argument_comparison_token( + context.arguments + ) argument_labels = [*input_labels, *resolved_labels] argument_label = combine_labels(*argument_labels) if argument_labels else ContentLabel() diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 1115252db87..f758eb6385d 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3,6 +3,8 @@ import asyncio import json import logging +import math +import threading import warnings from collections.abc import AsyncIterable, Awaitable, Callable, Sequence from typing import Any, Literal @@ -3231,6 +3233,151 @@ def count_items(count: int) -> str: assert validation_count == 1 +async def test_approved_coercing_arguments_execute_without_replacement() -> None: + """Approval binds to the normalized middleware-entry representation.""" + from agent_framework._tools import _auto_invoke_function, normalize_function_invocation_configuration + + class CountArgs(BaseModel): + count: int + + class PassthroughMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + assert context.arguments == {"count": 3} + await call_next() + + received: list[int] = [] + + @tool(name="approved_count", schema=CountArgs, approval_mode="always_require") + def approved_count(count: int) -> str: + received.append(count) + return str(count) + + function_call = Content.from_function_call( + call_id="approved-coercion", + id="approved-coercion-occurrence", + name=approved_count.name, + arguments={"count": "3"}, + ) + approval_response = Content.from_function_approval_request( + id="approved-coercion-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={approved_count.name: approved_count}, + middleware_pipeline=FunctionMiddlewarePipeline(PassthroughMiddleware()), + ) + + assert result.type == "function_result" + assert result.result == "3" + assert received == [3] + + +async def test_prepared_arguments_support_noncopyable_validator_output() -> None: + """Prepared argument reuse does not require validator outputs to be deepcopyable.""" + from agent_framework._tools import _auto_invoke_function, normalize_function_invocation_configuration + + validation_count = 0 + received: list[Any] = [] + + class LockArgs(BaseModel): + value: Any + + @field_validator("value") + @classmethod + def create_lock(cls, value: Any) -> Any: + nonlocal validation_count + validation_count += 1 + return threading.Lock() if value == "lock" else value + + class PassthroughMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + await call_next() + + @tool(name="lock_tool", schema=LockArgs, approval_mode="never_require") + def lock_tool(value: Any) -> str: + received.append(value) + return "locked" if value.locked() else "unlocked" + + result = await _auto_invoke_function( + Content.from_function_call(call_id="noncopyable", name=lock_tool.name, arguments={"value": "lock"}), + config=normalize_function_invocation_configuration(None), + tool_map={lock_tool.name: lock_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(PassthroughMiddleware()), + ) + + assert result.type == "function_result" + assert result.result == "unlocked" + assert len(received) == 1 + assert validation_count == 1 + + +async def test_nan_prepared_and_approval_snapshots_are_stable() -> None: + """An unchanged NaN survives prepared and approval snapshot comparison.""" + from agent_framework._tools import _auto_invoke_function, normalize_function_invocation_configuration + + validation_count = 0 + received: list[float] = [] + + class FloatArgs(BaseModel): + value: float + + @field_validator("value") + @classmethod + def track_validation(cls, value: float) -> float: + nonlocal validation_count + validation_count += 1 + return value + + class PassthroughMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + assert isinstance(context.arguments, dict) + assert math.isnan(context.arguments["value"]) + await call_next() + + @tool(name="nan_tool", schema=FloatArgs, approval_mode="always_require") + def nan_tool(value: float) -> str: + received.append(value) + return "nan" + + function_call = Content.from_function_call( + call_id="nan-call", + id="nan-occurrence", + name=nan_tool.name, + arguments={"value": "NaN"}, + ) + approval_response = Content.from_function_approval_request( + id="nan-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={nan_tool.name: nan_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(PassthroughMiddleware()), + ) + + assert result.type == "function_result" + assert result.result == "nan" + assert len(received) == 1 and math.isnan(received[0]) + assert validation_count == 1 + + async def test_function_middleware_can_short_circuit_before_argument_validation( chat_client_base: SupportsChatGetResponse, ) -> None: @@ -3538,8 +3685,16 @@ def approved_short_circuit(count: int) -> str: assert result.result == "handled by middleware" -async def test_approval_snapshot_distinguishes_boolean_from_integer() -> None: - """A boolean-to-integer mutation requires replacement approval.""" +@pytest.mark.parametrize( + ("approved_value", "changed_value"), + [(True, 1), (-0.0, 0.0)], + ids=["boolean-to-integer", "negative-zero-to-positive-zero"], +) +async def test_approval_snapshot_distinguishes_exact_values( + approved_value: Any, + changed_value: Any, +) -> None: + """A type or floating-point bit change requires replacement approval.""" from agent_framework._tools import _auto_invoke_function, normalize_function_invocation_configuration class ChangeTypeMiddleware(FunctionMiddleware): @@ -3548,7 +3703,7 @@ async def process( context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]], ) -> None: - context.arguments = {"value": 1} + context.arguments = {"value": changed_value} await call_next() @tool(name="typed_value", approval_mode="always_require") @@ -3559,7 +3714,7 @@ def typed_value(value: Any) -> str: call_id="typed-approval", id="typed-approval-occurrence", name=typed_value.name, - arguments={"value": True}, + arguments={"value": approved_value}, ) approval_response = Content.from_function_approval_request( id="typed-approval-occurrence", @@ -3578,7 +3733,14 @@ def typed_value(value: Any) -> str: assert isinstance(replacement, Content) assert replacement.type == "function_approval_request" assert replacement.function_call is not None - assert replacement.function_call.parse_arguments() == {"value": 1} + replacement_arguments = replacement.function_call.parse_arguments() + assert replacement_arguments is not None + replacement_value = replacement_arguments["value"] + if isinstance(changed_value, float): + assert isinstance(replacement_value, float) + assert math.copysign(1.0, replacement_value) == math.copysign(1.0, changed_value) + else: + assert replacement_value == changed_value async def test_hosted_tool_approval_response(chat_client_base: SupportsChatGetResponse): diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 16bdfad8618..b4932fa95ff 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -5,6 +5,7 @@ import asyncio import json import logging +import math from datetime import timedelta from types import SimpleNamespace from typing import Any, cast @@ -5913,21 +5914,30 @@ def reject_value(cls, value: str) -> str: assert "secret payload" not in str(result.exception) assert result.exception == "Invalid arguments for 'validator_type_error_sink'." - async def test_argument_mutation_after_security_middleware_fails_closed(self) -> None: - """A later middleware cannot change arguments after security policy has inspected them.""" + @pytest.mark.parametrize( + ("inspected_value", "changed_value"), + [(True, 1), (-0.0, 0.0)], + ids=["boolean-to-integer", "negative-zero-to-positive-zero"], + ) + async def test_argument_mutation_after_security_middleware_fails_closed( + self, + inspected_value: Any, + changed_value: Any, + ) -> None: + """A later type or float-bit change fails closed after policy inspection.""" tracker = LabelTrackingFunctionMiddleware() policy = PolicyEnforcementFunctionMiddleware() sink = self._sink(accepts_untrusted=True) class LateRepairMiddleware(FunctionMiddleware): async def process(self, context, call_next): - context.arguments = {"value": 1} + context.arguments = {"value": changed_value} await call_next() function_call = Content.from_function_call( call_id="late-repair", name=sink.name, - arguments={"value": True}, + arguments={"value": inspected_value}, ) with pytest.raises(MiddlewareFailure, match="Install argument-repair middleware before security middleware"): @@ -5938,6 +5948,41 @@ async def process(self, context, call_next): middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy, LateRepairMiddleware()), ) + async def test_security_snapshot_accepts_unchanged_nan(self) -> None: + """An unchanged NaN remains stable through security snapshot comparison.""" + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + received: list[float] = [] + + class FloatArgs(BaseModel): + value: float + + def nan_sink(value: float) -> str: + received.append(value) + return "nan" + + nan_tool = FunctionTool( + func=nan_sink, + name="nan_sink", + input_model=FloatArgs, + additional_properties={"accepts_untrusted": True}, + ) + function_call = Content.from_function_call( + call_id="security-nan", + name=nan_tool.name, + arguments={"value": "NaN"}, + ) + + result = await _auto_invoke_function( + function_call, + config=normalize_function_invocation_configuration(None), + tool_map={nan_tool.name: nan_tool}, + middleware_pipeline=FunctionMiddlewarePipeline(tracker, policy), + ) + + assert result.type == "function_result" + assert len(received) == 1 and math.isnan(received[0]) + async def test_security_policy_observes_custom_validator_transform_once(self) -> None: """Security middleware inspects the exact normalized value delivered to the tool.""" tracker = LabelTrackingFunctionMiddleware()