diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 4afdd0cd3bc..f16e263cbf8 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/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. diff --git a/python/packages/a2a/agent_framework_a2a/_utils.py b/python/packages/a2a/agent_framework_a2a/_utils.py index eba67957d36..64a3e6a7fa7 100644 --- a/python/packages/a2a/agent_framework_a2a/_utils.py +++ b/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[^;,]+(?:;[^;,=]+=[^;,]+)*);base64,(?P[A-Za-z0-9+/=]*)\Z") diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 21049245d2e..bc6c85e285e 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/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)}