Python: A2AAgent: forward selected run kwargs in request metadata

python/packages/a2a/agent_framework_a2a/_agent.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import base64
+import json
 import sys
 import uuid
 import warnings
@@ -48,9 +49,10 @@
 from agent_framework.exceptions import AgentInvalidRequestException
 from agent_framework.observability import AgentTelemetryLayer
 from google.protobuf.json_format import MessageToDict
+from google.protobuf.struct_pb2 import Struct
 
 from ._feature_usage import FeatureIndex
-from ._utils import get_uri_data
+from ._utils import AGENT_FRAMEWORK_METADATA_KEY, get_uri_data
 
 if sys.version_info >= (3, 11):
     from typing import TypedDict  # pragma: no cover
@@ -216,6 +218,48 @@ class A2AContinuationToken(ContinuationToken):
 A2AStreamItem: TypeAlias = StreamResponse
 
 
+# Largest magnitude at which every integer is exactly representable as the double protobuf Struct uses.
+_MAX_EXACT_STRUCT_INT = 2**53
+
+
+def _parse_exact_int(text: str) -> int:
+    value = int(text)
+    if abs(value) > _MAX_EXACT_STRUCT_INT:
+        raise ValueError(f"integer {text[:20]} is outside the exactly representable range of a protobuf Struct")
+    return value
+
+
+def _reject_non_str_keys(value: Any) -> None:
+    """Recursively raise if any mapping key is not a ``str``, since JSON would silently coerce it."""
+    if isinstance(value, Mapping):
+        for key, item in cast("Mapping[Any, Any]", value).items():
+            if not isinstance(key, str):
+                raise TypeError(f"mapping keys must be str, got {type(key).__name__}: {key!r}")
+            _reject_non_str_keys(item)
+    elif isinstance(value, (list, tuple)):
+        for item in cast("Sequence[Any]", value):
+            _reject_non_str_keys(item)
+
+
+def _to_struct_safe_json(values: Mapping[str, Any]) -> dict[str, Any]:
+    """Return a JSON round-tripped copy of ``values`` that a protobuf ``Struct`` can carry exactly.
+
+    Raises:
+        TypeError: If a value is not JSON-compatible or a mapping key is not a ``str``.
+        ValueError: If a value is NaN, infinite, an out-of-range integer, or otherwise unrepresentable.
+        OverflowError: If a number is too large for a protobuf ``Struct``.
+    """
+    # json.dumps coerces non-str keys (for example 1 and "1" collide), so reject them before the round trip.
+    _reject_non_str_keys(values)
+    # allow_nan=False rejects NaN and infinities, and parse_int rejects integers a Struct would round.
+    round_tripped = cast(
+        "dict[str, Any]", json.loads(json.dumps(dict(values), allow_nan=False), parse_int=_parse_exact_int)
+    )
+    # Dry run so any remaining Struct conversion failure (for example an unpaired surrogate) surfaces here.
+    Struct().update(round_tripped)
+    return round_tripped
+
+
 class A2AAgent(AgentTelemetryLayer, BaseAgent):
     """Agent2Agent (A2A) protocol implementation.
 
@@ -242,6 +286,7 @@ def __init__(
         auth_interceptor: AuthInterceptor | None = None,
         timeout: float | httpx.Timeout | None = None,
         supported_protocol_bindings: list[Literal["JSONRPC", "GRPC", "HTTP+JSON"] | str] | None = None,
+        forwarded_kwargs: Sequence[str] | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the A2AAgent.
@@ -264,6 +309,13 @@ def __init__(
             supported_protocol_bindings: List of protocol bindings to use for transport negotiation.
                 Known values: "JSONRPC", "GRPC", "HTTP+JSON". Defaults to ["JSONRPC"].
                 The A2A spec treats this as an open-form string, so custom bindings are also accepted.
+            forwarded_kwargs: Names of keys to forward to the remote A2A server. Keys are looked up in
+                the ``function_invocation_kwargs`` and ``client_kwargs`` passed to :meth:`run` and sent
+                in ``SendMessageRequest.metadata`` under ``"agent_framework"``. The default (``None``)
+                sends nothing. Values must be JSON-compatible with finite numbers and integers within
+                +/-2**53 (protobuf ``Struct`` stores numbers as doubles), otherwise
+                ``AgentInvalidRequestException`` is raised. The receiving ``A2AExecutor`` decides
+                which keys it accepts.
             kwargs: any additional properties, passed to BaseAgent.
         """
         # Default name/description from agent_card when not explicitly provided
@@ -274,6 +326,7 @@ def __init__(
                 description = agent_card.description
 
         super().__init__(id=id, name=name, description=description, **kwargs)
+        self._forwarded_kwargs: frozenset[str] = frozenset(forwarded_kwargs or ())
         self._http_client: httpx.AsyncClient | None = http_client
         # Only HTTP clients created by this agent are closed on exit.
         self._close_http_client = False
@@ -490,10 +543,10 @@ def run(
         Keyword Args:
             stream: Whether to stream the response. Defaults to False.
             session: The conversation session associated with the message(s).
-            function_invocation_kwargs: Present for compatibility with the shared agent interface.
-                A2AAgent does not use these values directly.
-            client_kwargs: Present for compatibility with the shared agent interface.
-                A2AAgent does not use these values directly.
+            function_invocation_kwargs: Keyword arguments for tool invocation. Only the keys named in
+                ``forwarded_kwargs`` are sent to the server, nothing otherwise.
+            client_kwargs: Keyword arguments for the chat client. Only the keys named in
+                ``forwarded_kwargs`` are sent to the server, nothing otherwise.
             kwargs: Additional compatibility keyword arguments.
                 A2AAgent does not use these values directly.
             continuation_token: Optional token to resume a long-running task
@@ -535,6 +588,9 @@ def run(
             a2a_message = self._prepare_message_for_a2a(normalized_messages[-1], session=session)
             input_request_occurrence_id = a2a_message.message_id
             request = SendMessageRequest(message=a2a_message)
+            forwarded = self._get_forwarded_kwargs(function_invocation_kwargs, client_kwargs)
+            if forwarded:
+                request.metadata.update({AGENT_FRAMEWORK_METADATA_KEY: forwarded})
             if background and not stream:
                 # return_immediately only applies to non-streaming (message/send)
                 request.configuration.return_immediately = True
@@ -1089,6 +1145,35 @@ async def poll_task(self, continuation_token: A2AContinuationToken) -> AgentResp
             return AgentResponse.from_updates(updates)
         return AgentResponse(messages=[], response_id=task.id, raw_representation=task)
 
+    def _get_forwarded_kwargs(
+        self,
+        function_invocation_kwargs: Mapping[str, Any] | None,
+        client_kwargs: Mapping[str, Any] | None,
+    ) -> dict[str, dict[str, Any]]:
+        """Select the opted-in run kwargs to send in the request metadata.
+
+        Raises:
+            AgentInvalidRequestException: If a selected value is not JSON-compatible.
+        """
+        if not self._forwarded_kwargs:
+            return {}
+        forwarded: dict[str, dict[str, Any]] = {}
+        for name, source in (
+            ("function_invocation_kwargs", function_invocation_kwargs),
+            ("client_kwargs", client_kwargs),
+        ):
+            selected = {key: value for key, value in (source or {}).items() if key in self._forwarded_kwargs}
+            if not selected:
+                continue
+            try:
+                forwarded[name] = _to_struct_safe_json(selected)
+            except (TypeError, ValueError, OverflowError) as ex:
+                raise AgentInvalidRequestException(
+                    f"Forwarded {name} must contain only values that protobuf Struct can carry "
+                    f"(JSON-compatible, finite numbers, integers within +/-2**53): {ex}"
+                ) from ex
+        return forwarded
+
     def _prepare_message_for_a2a(self, message: Message, *, session: AgentSession | None = None) -> A2AMessage:
         """Prepare a Message for the A2A protocol.
 

python/packages/a2a/agent_framework_a2a/_utils.py

@@ -2,6 +2,10 @@
 
 import re
 
+# Key in ``SendMessageRequest.metadata`` that carries forwarded ``function_invocation_kwargs`` and
+# ``client_kwargs``.
+AGENT_FRAMEWORK_METADATA_KEY = "agent_framework"
+
 URI_PATTERN = re.compile(r"^data:(?P<media_type>[^;,]+(?:;[^;,=]+=[^;,]+)*);base64,(?P<base64_data>[A-Za-z0-9+/=]*)\Z")
 
 

python/packages/a2a/tests/test_a2a_agent.py

@@ -2896,3 +2896,97 @@ async def test_context_manager_closes_self_created_http_client() -> None:
             pass
 
         mock_http_client.aclose.assert_called_once()
+
+
+def _forwarding_agent(mock_a2a_client: MockA2AClient, forwarded_kwargs: list[str] | None) -> A2AAgent:
+    return A2AAgent(
+        name="Test Agent",
+        id="test-agent",
+        client=cast(Any, mock_a2a_client),
+        http_client=None,
+        forwarded_kwargs=forwarded_kwargs,
+    )
+
+
+async def test_run_forwards_selected_kwargs_in_request_metadata(mock_a2a_client: MockA2AClient) -> None:
+    """Only keys named in forwarded_kwargs are sent, grouped by their source kwargs, in the request metadata."""
+    agent = _forwarding_agent(mock_a2a_client, ["tenant", "trace"])
+    mock_a2a_client.add_message_response("msg-1", "Done")
+
+    await agent.run(
+        "Hello",
+        function_invocation_kwargs={"tenant": "acme", "secret": "do-not-send"},
+        client_kwargs={"trace": {"id": "abc"}, "other": 1},
+    )
+
+    assert MessageToDict(mock_a2a_client.last_request.metadata) == {
+        "agent_framework": {
+            "function_invocation_kwargs": {"tenant": "acme"},
+            "client_kwargs": {"trace": {"id": "abc"}},
+        }
+    }
+    assert not mock_a2a_client.last_message.metadata
+
+
+async def test_run_forwards_nothing_by_default(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
+    mock_a2a_client.add_message_response("msg-1", "Done")
+
+    await a2a_agent.run("Hello", function_invocation_kwargs={"tenant": "acme"}, client_kwargs={"trace": "abc"})
+
+    assert not mock_a2a_client.last_request.metadata
+
+
+async def test_run_forwards_nothing_when_selected_keys_absent(mock_a2a_client: MockA2AClient) -> None:
+    agent = _forwarding_agent(mock_a2a_client, ["tenant"])
+    mock_a2a_client.add_message_response("msg-1", "Done")
+
+    await agent.run("Hello", client_kwargs={"trace": "abc"})
+
+    assert not mock_a2a_client.last_request.metadata
+
+
+async def test_run_rejects_non_json_forwarded_values(mock_a2a_client: MockA2AClient) -> None:
+    agent = _forwarding_agent(mock_a2a_client, ["tenant"])
+
+    with raises(AgentInvalidRequestException):
+        await agent.run("Hello", function_invocation_kwargs={"tenant": object()})
+
+    assert mock_a2a_client.call_count == 0
+
+
+@mark.parametrize(
+    "value",
+    [float("nan"), float("inf"), -float("inf"), 2**53 + 1, -(2**53) - 1, 10**400, "\ud800"],
+    ids=["nan", "inf", "-inf", "int-too-large", "int-too-small", "int-overflow", "unpaired-surrogate"],
+)
+async def test_run_rejects_values_protobuf_struct_cannot_carry(mock_a2a_client: MockA2AClient, value: Any) -> None:
+    agent = _forwarding_agent(mock_a2a_client, ["tenant"])
+
+    with raises(AgentInvalidRequestException):
+        await agent.run("Hello", function_invocation_kwargs={"tenant": value})
+
+    assert mock_a2a_client.call_count == 0
+
+
+@mark.parametrize(
+    "value",
+    [{1: "numeric", "1": "string"}, {"nested": [{2: "x"}]}, {None: "x"}],
+    ids=["key-collision", "nested-in-list", "none-key"],
+)
+async def test_run_rejects_non_str_mapping_keys(mock_a2a_client: MockA2AClient, value: Any) -> None:
+    agent = _forwarding_agent(mock_a2a_client, ["tenant"])
+
+    with raises(AgentInvalidRequestException):
+        await agent.run("Hello", function_invocation_kwargs={"tenant": value})
+
+    assert mock_a2a_client.call_count == 0
+
+
+async def test_run_forwards_integers_at_the_exact_boundary(mock_a2a_client: MockA2AClient) -> None:
+    agent = _forwarding_agent(mock_a2a_client, ["count"])
+    mock_a2a_client.add_message_response("msg-1", "Done")
+
+    await agent.run("Hello", client_kwargs={"count": 2**53})
+
+    forwarded = MessageToDict(mock_a2a_client.last_request.metadata)["agent_framework"]["client_kwargs"]
+    assert forwarded == {"count": float(2**53)}