diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index eb40fec7ee7..626855b7c8a 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/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, diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index 8dd1da6f053..6d7a381ff5a 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/python/packages/core/tests/core/test_tools.py @@ -3,9 +3,10 @@ import copy import logging import threading +from collections.abc import Awaitable, Callable from contextvars import ContextVar from datetime import date, datetime -from typing import Annotated, Any, Literal, get_args, get_origin +from typing import Annotated, Any, Literal, cast, get_args, get_origin from unittest.mock import Mock import pytest @@ -20,7 +21,11 @@ FunctionTool, tool, ) -from agent_framework._middleware import FunctionInvocationContext +from agent_framework._middleware import ( + FunctionInvocationContext, + FunctionMiddleware, + FunctionMiddlewarePipeline, +) from agent_framework._tools import ( _auto_invoke_function, _format_tool_parameters, @@ -431,6 +436,291 @@ def search(query: str) -> str: await search.invoke(arguments={}) +@pytest.mark.parametrize("arguments", ["not {json", "[1, 2]", '"text"', " "]) +async def test_function_call_arguments_that_are_not_a_json_object_are_answered_with_an_error( + arguments: str, +) -> None: + """A payload that is not a JSON object is not arguments, and must not be called as one. + + Wrapping it under "raw" would call the tool with an argument it never + declared, and the model would not learn that its call was malformed. + """ + calls: list[dict[str, Any]] = [] + + async def list_accounts(**kwargs: Any) -> str: + calls.append(kwargs) + return "ok" + + account_tool = FunctionTool( + name="list_accounts", + description="List the accounts", + func=list_accounts, + input_model={"type": "object", "properties": {}}, + ) + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments=arguments), + ], + tools=[account_tool], + config={}, + ) + + assert not should_terminate + assert calls == [] + result = result_groups[0][0] + assert "not a JSON object" in str(result.exception) + + +async def test_function_call_with_object_arguments_still_runs_the_tool() -> None: + """A well-formed call is unaffected by the guard.""" + calls: list[dict[str, Any]] = [] + + async def list_accounts(**kwargs: Any) -> str: + calls.append(kwargs) + return "ok" + + account_tool = FunctionTool( + name="list_accounts", + description="List the accounts", + func=list_accounts, + input_model={"type": "object", "properties": {}}, + ) + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments='{"limit": 5}'), + ], + tools=[account_tool], + config={}, + ) + + assert not should_terminate + assert calls == [{"limit": 5}] + assert result_groups[0][0].result == "ok" + + +async def test_malformed_arguments_reach_approval_like_any_other_call() -> None: + """A malformed call is not answered before approval. + + Middleware may still repair the arguments once the call is approved, so the + request goes to approval the way every other call does; the innermost + handler answers the call only if the payload is still unrepaired there. + """ + calls: list[dict[str, Any]] = [] + + @tool(name="guarded_write", approval_mode="always_require") + def guarded_write(value: str = "") -> str: + calls.append({"value": value}) + return "ok" + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="guarded_write", arguments="not {json"), + ], + tools=[guarded_write], + config={}, + ) + + assert not should_terminate + assert calls == [] + assert len(result_groups) == 1 + assert result_groups[0][0].type == "function_approval_request" + + +async def test_middleware_can_repair_a_payload_that_is_not_a_json_object() -> None: + """Middleware sees the raw parsed mapping and may repair it into real arguments. + + This is the function-loop contract: middleware receives the arguments before + any final validation, and a repaired call executes like any other. + """ + + calls: list[dict[str, Any]] = [] + + @tool(name="list_accounts") + def list_accounts(limit: int = 0) -> str: + calls.append({"limit": limit}) + return "ok" + + class RepairRawArguments(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + arguments = cast("dict[str, Any]", context.arguments) + if "raw" in arguments: + context.arguments = {"limit": 5} + await call_next() + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments="[1, 2]"), + ], + tools=[list_accounts], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(RepairRawArguments()), + ) + + assert not should_terminate + assert calls == [{"limit": 5}] + assert result_groups[0][0].result == "ok" + + +@pytest.mark.parametrize("arguments", ["not {json", "[1, 2]"]) +async def test_malformed_arguments_left_unrepaired_are_answered_before_the_tool_runs( + arguments: str, +) -> None: + """When middleware does not repair the wrapped payload, the call is answered. + + The error is the argument-validation result, raised immediately before the + tool would have executed, so the tool never runs on undeclared arguments. + """ + calls: list[dict[str, Any]] = [] + + async def list_accounts(**kwargs: Any) -> str: + calls.append(kwargs) + return "ok" + + account_tool = FunctionTool( + name="list_accounts", + description="List the accounts", + func=list_accounts, + input_model={"type": "object", "properties": {}}, + ) + + class Passthrough(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + await call_next() + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments=arguments), + ], + tools=[account_tool], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(Passthrough()), + ) + + assert not should_terminate + assert calls == [] + result = result_groups[0][0] + assert "not a JSON object" in str(result.exception) + + +class _InjectUser(FunctionMiddleware): + """Middleware that appends a key to every call's arguments.""" + + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.arguments = {**cast("dict[str, Any]", context.arguments), "user": "u1"} + await call_next() + + +async def test_injecting_middleware_does_not_make_a_malformed_payload_count_as_repaired() -> None: + """A middleware that appends a key must not un-mark the wrapped payload. + + The injection preserves the original ``raw`` entry, so the arguments are + still the wrapper plus an undeclared key - the tool must not run. + """ + calls: list[dict[str, Any]] = [] + + async def list_accounts(**kwargs: Any) -> str: + calls.append(kwargs) + return "ok" + + account_tool = FunctionTool( + name="list_accounts", + description="List the accounts", + func=list_accounts, + input_model={"type": "object", "properties": {}}, + ) + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments="[1, 2]"), + ], + tools=[account_tool], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(_InjectUser()), + ) + + assert not should_terminate + assert calls == [] + result = result_groups[0][0] + assert "not a JSON object" in str(result.exception) + + +async def test_injecting_middleware_does_not_silence_a_malformed_payload_for_a_typed_tool() -> None: + """The same holds for a typed tool: the wrapper must not reach it either. + + The input model would silently drop the undeclared ``raw`` key, so the tool + would run on injected values alone without the model ever learning that its + call was malformed. + """ + calls: list[dict[str, Any]] = [] + + @tool(name="typed_target") + def typed_target(user: str = "", limit: int = 0) -> str: + calls.append({"user": user, "limit": limit}) + return "ok" + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="typed_target", arguments="[1, 2]"), + ], + tools=[typed_target], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(_InjectUser()), + ) + + assert not should_terminate + assert calls == [] + result = result_groups[0][0] + assert "not a JSON object" in str(result.exception) + + +async def test_a_tool_that_declares_raw_receives_a_non_object_payload_as_arguments() -> None: + """A tool that declares ``raw`` gets the wrapped payload as ordinary arguments. + + For such a tool the wrapper coincides with a valid call, which is the + behaviour the loop had before the malformed-argument guard existed. + """ + calls: list[dict[str, Any]] = [] + + @tool(name="raw_writer") + def raw_writer(raw: str = "") -> str: + calls.append({"raw": raw}) + return "ok" + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="raw_writer", arguments="not {json"), + ], + tools=[raw_writer], + config={}, + ) + + assert not should_terminate + assert calls == [{"raw": "not {json"}] + assert result_groups[0][0].result == "ok" + + async def test_invoke_preserves_explicit_null_argument(): """A required nullable argument the model sets to null must reach the function. @@ -2168,3 +2458,196 @@ def test_skip_parsing_is_singleton() -> None: # endregion + + +@pytest.mark.parametrize("arguments", ["null", "NaN"]) +async def test_malformed_null_and_nan_payloads_are_answered_without_middleware(arguments: str) -> None: + """A wrapper whose raw entry is None or NaN is still a wrapper. + + ``parse_arguments()`` wraps ``"null"`` as ``{"raw": None}`` and ``"NaN"`` + as ``{"raw": }``; neither the absence of a value nor the NaN + inequality may un-mark the wrapper. + """ + calls: list[dict[str, Any]] = [] + + async def list_accounts(**kwargs: Any) -> str: + calls.append(kwargs) + return "ok" + + account_tool = FunctionTool( + name="list_accounts", + description="List the accounts", + func=list_accounts, + input_model={"type": "object", "properties": {}}, + ) + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments=arguments), + ], + tools=[account_tool], + config={}, + ) + + assert not should_terminate + assert calls == [] + result = result_groups[0][0] + assert "not a JSON object" in str(result.exception) + + +@pytest.mark.parametrize("arguments", ["null", "NaN"]) +async def test_injecting_middleware_does_not_unmark_null_or_nan_wrappers(arguments: str) -> None: + """An injecting middleware keeps the raw entry, so the call stays malformed. + + The injected key changes the mapping but the wrapper survives with its + original None/NaN value, and the wrapper-presence check is value-safe for + both. + """ + calls: list[dict[str, Any]] = [] + + async def list_accounts(**kwargs: Any) -> str: + calls.append(kwargs) + return "ok" + + account_tool = FunctionTool( + name="list_accounts", + description="List the accounts", + func=list_accounts, + input_model={"type": "object", "properties": {}}, + ) + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments=arguments), + ], + tools=[account_tool], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(_InjectUser()), + ) + + assert not should_terminate + assert calls == [] + result = result_groups[0][0] + assert "not a JSON object" in str(result.exception) + + +class _ReplaceWithValidArguments(FunctionMiddleware): + """Middleware that replaces the arguments with a valid mapping. + + The replacement drops any raw wrapper and supplies declared arguments, + which is a full repair. + """ + + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.arguments = {"limit": 5} + await call_next() + + +async def test_middleware_that_replaces_null_wrapper_with_valid_arguments_repairs_the_call() -> None: + """A full-argument replacement is a repair, not an unrepaired wrapper. + + The replacement no longer contains the ``raw`` entry, so the + ``None``-valued wrapper of ``"null"`` must not be misread from the + absent key, and the repaired call executes. + """ + calls: list[dict[str, Any]] = [] + + @tool(name="list_accounts") + def list_accounts(limit: int = 0) -> str: + calls.append({"limit": limit}) + return "ok" + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="list_accounts", arguments="null"), + ], + tools=[list_accounts], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(_ReplaceWithValidArguments()), + ) + + assert not should_terminate + assert calls == [{"limit": 5}] + assert result_groups[0][0].result == "ok" + + +async def test_builtin_label_tracking_middleware_preserves_malformed_wrapper_provenance() -> None: + """The built-in label-tracking middleware must not erase the wrapper. + + Its argument normalization validates against the declared model and drops + the undeclared ``raw`` entry for a typed tool; the provenance marker keeps + the malformed payload rejected instead of executing with defaults. + """ + from agent_framework.security import LabelTrackingFunctionMiddleware + + calls: list[dict[str, Any]] = [] + + @tool(name="typed_label_target") + def typed_label_target(user: str = "", limit: int = 0) -> str: + calls.append({"user": user, "limit": limit}) + return "ok" + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call(call_id="call-1", name="typed_label_target", arguments="not {json"), + ], + tools=[typed_label_target], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(LabelTrackingFunctionMiddleware()), + ) + + assert not should_terminate + assert calls == [] + result = result_groups[0][0] + assert result.type == "function_result" + assert "Invalid arguments for 'typed_label_target'" in str(result.exception) + + +@tool(name="raw_reader", description="Read a raw value") +def raw_reader_tool(raw: datetime) -> str: + return f"raw:{raw.isoformat()}" + + +async def test_declared_raw_tool_keeps_normal_preparation_types() -> None: + """A declared ``raw: datetime`` reaches middleware as a datetime, not a string. + + Preparation runs normally for tools that declare ``raw``, so middleware + sees the declared types instead of the raw JSON string. + """ + seen_types: list[str] = [] + + class TrackTypes(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + raw_value = cast("dict[str, Any]", context.arguments).get("raw") + seen_types.append(type(raw_value).__name__) + await call_next() + + result_groups, should_terminate = await _try_execute_function_call_groups( + custom_args={}, + function_calls=[ + Content.from_function_call( + call_id="call-1", + name="raw_reader", + arguments='"2024-01-01T00:00:00+00:00"', + ), + ], + tools=[raw_reader_tool], + config={}, + middleware_pipeline=FunctionMiddlewarePipeline(TrackTypes()), + ) + + assert not should_terminate + assert seen_types == ["datetime"] + assert result_groups[0][0].result == "raw:2024-01-01T00:00:00+00:00"