Python: fix(tools): answer a function call whose arguments are not a JSON object

python/packages/core/agent_framework/_tools.py

@@ -136,6 +136,7 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None)
 _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"
+_MALFORMED_WRAPPER_CONTEXT_KEY: Final[str] = "_malformed_wrapper_provenance"
 _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."
@@ -1020,7 +1021,18 @@ def _prepare_context_arguments(
             context.arguments = prepared_state.arguments
             return prepared_state.arguments
 
+        # A malformed payload provenance marker keeps the synthetic raw entry
+        # alive across schema validation: validation drops the undeclared key,
+        # which would otherwise erase the signal that the call is still the
+        # wrapped malformed payload rather than a middleware repair.
+        preserved_raw = None
+        malformed_wrapper = context.metadata.get(_MALFORMED_WRAPPER_CONTEXT_KEY)
+        if malformed_wrapper is not None and isinstance(current_arguments, Mapping) and "raw" in current_arguments:
+            preserved_raw = current_arguments["raw"]
+
         validated_arguments = self._prepare_arguments(arguments)
+        if preserved_raw is not None:
+            validated_arguments = {**validated_arguments, "raw": preserved_raw}
         context.arguments = validated_arguments
         context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = _PreparedArgumentsState(
             arguments=validated_arguments,
@@ -2164,6 +2176,40 @@ def _finalize_function_result(
     return function_result
 
 
+def _tool_declares_parameter(tool: FunctionTool, name: str) -> bool:
+    """Return whether the tool's own schema declares a parameter with this name.
+
+    Schema-supplied tools declare through their JSON ``properties``; decorated
+    tools through their resolved input model. A ``**kwargs`` catch-all does not
+    declare anything: the model never agreed to such a parameter by name.
+    """
+    if tool._schema_supplied:  # pyright: ignore[reportPrivateUsage]
+        properties = tool._input_schema_cached.get("properties") if tool._input_schema_cached else None  # pyright: ignore[reportPrivateUsage]
+        return isinstance(properties, Mapping) and name in properties
+    return tool.input_model is not None and name in tool.input_model.model_fields
+
+
+def _function_call_payload_is_not_an_object(function_call: Content) -> bool:
+    """Return whether the model sent a payload that is not a JSON object.
+
+    `Content.parse_arguments()` wraps such a payload under a single "raw" key.
+    That is the right answer for a caller that wants to see what the model sent,
+    but not for a function call: the tool would run with an argument it never
+    declared, and the model would not learn that its call was malformed. The
+    wrapped mapping still reaches the function middleware, which can repair the
+    arguments the way it repairs any other invalid arguments; when it stays
+    unrepaired, the innermost handler answers the call instead of running it.
+    """
+    arguments = function_call.arguments
+    if isinstance(arguments, str) and arguments:
+        try:
+            decoded = json.loads(arguments)
+        except (json.JSONDecodeError, TypeError):
+            decoded = None
+        return not isinstance(decoded, dict)
+    return False
+
+
 async def _auto_invoke_function(
     function_call_content: Content,
     custom_args: dict[str, Any] | None = None,
@@ -2239,7 +2285,37 @@ async def _auto_invoke_function(
         approval_response = function_call_content
         function_call_content = approved_function_call
 
-    parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
+    # A payload that is not a JSON object still reaches middleware as the raw
+    # parsed mapping it wraps, so middleware has the same chance to repair it
+    # that it has for any other invalid arguments. If it is still unrepaired,
+    # the innermost handler answers it before the tool runs.
+    payload_is_not_an_object = _function_call_payload_is_not_an_object(function_call_content)
+    parsed_args = dict(function_call_content.parse_arguments() or {})
+    wrapped_raw_present = payload_is_not_an_object and set(parsed_args) == {"raw"}
+    wrapped_raw_value = parsed_args.get("raw") if wrapped_raw_present else None
+
+    def _unrepaired_wrapper(arguments: Any) -> bool:
+        """Return whether the arguments are still the wrapped malformed payload.
+
+        Wrapper presence is tracked separately from the wrapped value: ``None``
+        and ``NaN`` are legitimate contents of a wrapper, not signs that one is
+        absent. The wrapper counts as repaired only when middleware replaced
+        it. Keeping the original ``raw`` entry - which an argument-injecting
+        middleware does - leaves the call malformed, while a ``raw`` entry the
+        tool itself declares is ordinary arguments and never counts as the
+        wrapper.
+        """
+        if not wrapped_raw_present or _tool_declares_parameter(tool, "raw"):
+            return False
+        if not isinstance(arguments, Mapping) or "raw" not in arguments:
+            return False
+        raw_value = cast("Mapping[str, Any]", arguments)["raw"]
+        return raw_value == wrapped_raw_value or (
+            isinstance(raw_value, float)
+            and isinstance(wrapped_raw_value, float)
+            and raw_value != raw_value
+            and wrapped_raw_value != wrapped_raw_value
+        )
 
     # Filter out internal framework kwargs before passing to tools.
     # conversation_id is an internal tracking ID that should not be forwarded to tools.
@@ -2278,6 +2354,15 @@ async def _auto_invoke_function(
         # No middleware - execute directly
         direct_context = None
         try:
+            if _unrepaired_wrapper(args):
+                # Without middleware there is nothing that could repair the
+                # wrapped payload, so answer it before the tool runs. A tool
+                # that declares ``raw`` receives the payload as ordinary
+                # arguments instead.
+                raise _FunctionArgumentValidationError(
+                    f"Invalid arguments for '{tool.name}': arguments are not a JSON object",
+                    redacted_message=f"Invalid arguments for '{tool.name}'.",
+                )
             if getattr(tool, "_context_parameter_name", None):
                 direct_context = FunctionInvocationContext(
                     function=tool,
@@ -2312,13 +2397,18 @@ async def _auto_invoke_function(
             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
+    if not payload_is_not_an_object or _tool_declares_parameter(tool, "raw"):
+        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
+    # Only a wrapped malformed payload that the tool does not accept as a
+    # declared ``raw`` argument skips preparation: preparation would silently
+    # drop that wrapper as an extra field, and middleware would never see what
+    # the model actually sent.
 
     middleware_context = FunctionInvocationContext(
         function=tool,
@@ -2332,6 +2422,8 @@ 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 wrapped_raw_present:
+        middleware_context.metadata[_MALFORMED_WRAPPER_CONTEXT_KEY] = dict(parsed_args)
     if arguments_prepared:
         middleware_context.metadata[_PREPARED_ARGUMENTS_CONTEXT_KEY] = _PreparedArgumentsState(
             arguments=args,
@@ -2361,6 +2453,15 @@ async def _auto_invoke_function(
     async def final_function_handler(context_obj: Any) -> Any:
         nonlocal final_handler_started
         final_handler_started = True
+        if _unrepaired_wrapper(context_obj.arguments):
+            # Middleware left the wrapped payload in place - an argument-injecting
+            # middleware keeps its ``raw`` entry with the original value - so the
+            # arguments are still not arguments: answer it here, immediately
+            # before the tool runs.
+            raise _FunctionArgumentValidationError(
+                f"Invalid arguments for '{tool.name}': arguments are not a JSON object",
+                redacted_message=f"Invalid arguments for '{tool.name}'.",
+            )
         return await tool.invoke(
             arguments=context_obj.arguments,
             context=context_obj,