diff --git a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py index 2aa028d3856..1958b73ac28 100644 --- a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py +++ b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py @@ -4,9 +4,9 @@ import logging import uuid from asyncio import CancelledError -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from functools import partial -from typing import Any +from typing import Any, cast from a2a.helpers import new_task_from_user_message from a2a.server.agent_execution import AgentExecutor, RequestContext @@ -27,6 +27,22 @@ logger = logging.getLogger("agent_framework.a2a") +# Key in ``SendMessageRequest.metadata`` that carries forwarded ``function_invocation_kwargs`` and ``client_kwargs``. +AGENT_FRAMEWORK_METADATA_KEY = "agent_framework" +# Names owned by the framework (chat client and tool invocation parameters). A remote caller must never be able to +# set them, since a duplicate keyword would fail the task or override framework behavior. +_RESERVED_KWARGS = frozenset({ + "messages", + "stream", + "options", + "session", + "middleware", + "compaction_strategy", + "tokenizer", + "function_invocation_kwargs", + "client_kwargs", +}) + class A2AExecutor(AgentExecutor): """Execute AI agents using the A2A (Agent-to-Agent) protocol. @@ -89,9 +105,23 @@ class A2AExecutor(AgentExecutor): agent: The AI agent to execute. stream: Whether to stream the agent response. Defaults to False. run_kwargs: Additional keyword arguments to pass to the agent's run method. + accepted_function_invocation_kwargs: Names of keys the executor accepts from the caller's + forwarded ``function_invocation_kwargs`` (see ``A2AAgent``). A client sends them in the + request metadata, and only the keys named here are passed to the agent's run method. + The default (``None``) accepts nothing. + accepted_client_kwargs: Names of keys the executor accepts from the caller's forwarded + ``client_kwargs``. Authorized independently of ``accepted_function_invocation_kwargs``, + because these values reach the chat client. The default (``None``) accepts nothing. """ - def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Mapping[str, Any] | None = None): + def __init__( + self, + agent: SupportsAgentRun, + stream: bool = False, + run_kwargs: Mapping[str, Any] | None = None, + accepted_function_invocation_kwargs: Sequence[str] | None = None, + accepted_client_kwargs: Sequence[str] | None = None, + ): """Initialize the A2AExecutor with the specified agent. Args: @@ -99,9 +129,21 @@ def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Ma stream: Whether to stream the agent response. Defaults to False. run_kwargs: Additional keyword arguments to pass to the agent's run method. Cannot contain 'session' or 'stream' as these are managed by the executor. + accepted_function_invocation_kwargs: Names of keys to accept from the request's forwarded + ``function_invocation_kwargs``. Accepted values are passed to the agent's run method. + Values configured in ``run_kwargs`` take precedence on conflict. Defaults to accepting + nothing. + accepted_client_kwargs: Names of keys to accept from the request's forwarded + ``client_kwargs``. Authorized separately from ``accepted_function_invocation_kwargs``, + so opting into a tool-context key does not also let the caller inject it into the + chat client. Defaults to accepting nothing. + + Framework-owned names (for example ``stream``, ``messages``, ``options`` and ``session``) and + names starting with an underscore are rejected in both allowlists. Raises: - ValueError: If run_kwargs contains 'session' or 'stream'. + ValueError: If run_kwargs contains 'session' or 'stream', or an accepted kwargs allowlist + contains a framework-owned name. """ super().__init__() self._agent: SupportsAgentRun = agent @@ -112,6 +154,17 @@ def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Ma if "stream" in run_kwargs: raise ValueError("run_kwargs cannot contain 'stream' as it is managed by the executor.") self._run_kwargs: Mapping[str, Any] = run_kwargs or {} + for allowlist_name, allowlist in ( + ("accepted_function_invocation_kwargs", accepted_function_invocation_kwargs), + ("accepted_client_kwargs", accepted_client_kwargs), + ): + reserved = sorted(key for key in allowlist or () if key in _RESERVED_KWARGS or key.startswith("_")) + if reserved: + raise ValueError(f"{allowlist_name} cannot contain framework-owned names: {', '.join(reserved)}.") + self._accepted_kwargs: dict[str, frozenset[str]] = { + "function_invocation_kwargs": frozenset(accepted_function_invocation_kwargs or ()), + "client_kwargs": frozenset(accepted_client_kwargs or ()), + } @override async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None: @@ -165,10 +218,11 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non session = self._agent.create_session(session_id=task.context_id) + run_kwargs = self._build_run_kwargs(context) if self._stream: - await self._run_stream(query, session, updater) + await self._run_stream(query, session, updater, run_kwargs) else: - await self._run(query, session, updater) + await self._run(query, session, updater, run_kwargs) # Mark as complete await updater.complete() @@ -181,9 +235,31 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non message=updater.new_agent_message([Part(text=str(e))]), ) - async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None: + def _build_run_kwargs(self, context: RequestContext) -> dict[str, Any]: + """Merge the accepted request kwargs under the configured ``run_kwargs``.""" + run_kwargs: dict[str, Any] = dict(self._run_kwargs) + if not any(self._accepted_kwargs.values()): + return run_kwargs + forwarded = context.metadata.get(AGENT_FRAMEWORK_METADATA_KEY) + if not isinstance(forwarded, Mapping): + return run_kwargs + for name, allowed in self._accepted_kwargs.items(): + values = cast("Mapping[str, Any]", forwarded).get(name) + if not isinstance(values, Mapping): + continue + accepted = {key: value for key, value in cast("Mapping[str, Any]", values).items() if key in allowed} + if accepted: + configured = run_kwargs.get(name) + run_kwargs[name] = {**accepted, **(configured or {})} + return run_kwargs + + async def _run_stream( + self, query: Any, session: AgentSession, updater: TaskUpdater, run_kwargs: Mapping[str, Any] | None = None + ) -> None: """Run the agent in streaming mode and publish updates to the task updater.""" - response_stream = self._agent.run(query, session=session, stream=True, **self._run_kwargs) + response_stream = self._agent.run( + query, session=session, stream=True, **(self._run_kwargs if run_kwargs is None else run_kwargs) + ) streamed_artifact_ids: set[str] = set() # Generate a stable artifact ID for the entire stream so all chunks share the same ID. # This ensures clients can coalesce streaming tokens into a single artifact/message @@ -200,9 +276,13 @@ async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpda ) ).get_final_response() - async def _run(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None: + async def _run( + self, query: Any, session: AgentSession, updater: TaskUpdater, run_kwargs: Mapping[str, Any] | None = None + ) -> None: """Run the agent in non-streaming mode and publish messages to the task updater.""" - response = await self._agent.run(query, session=session, stream=False, **self._run_kwargs) + response = await self._agent.run( + query, session=session, stream=False, **(self._run_kwargs if run_kwargs is None else run_kwargs) + ) response_messages = response.messages if not isinstance(response_messages, list): diff --git a/python/packages/a2a/tests/test_a2a_executor.py b/python/packages/a2a/tests/test_a2a_executor.py index b0929193b12..a215a853de5 100644 --- a/python/packages/a2a/tests/test_a2a_executor.py +++ b/python/packages/a2a/tests/test_a2a_executor.py @@ -928,3 +928,155 @@ async def test_full_execution_flow_with_responses( mock_updater.start_work.assert_called_once() cast(Any, executor.handle_events).assert_called_once() mock_updater.complete.assert_called_once() + + +class TestA2AExecutorForwardedKwargs: + """Tests for accepting forwarded function_invocation_kwargs and client_kwargs.""" + + @staticmethod + async def _execute(executor: A2AExecutor, metadata: Any, mock_event_queue: MagicMock, mock_task: Task) -> MagicMock: + request_context = MagicMock() + request_context.get_user_input = MagicMock(return_value="Hello") + request_context.current_task = mock_task + request_context.context_id = "ctx-123" + request_context.message = MagicMock() + request_context.metadata = metadata + + response = MagicMock(spec=AgentResponse) + response.messages = [Message(role="assistant", contents=[Content.from_text(text="Hi")])] + cast(Any, executor._agent).run = AsyncMock(return_value=response) + cast(Any, executor._agent).create_session = MagicMock() + + with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class: + mock_updater = MagicMock() + for name in ("submit", "start_work", "complete", "update_status"): + setattr(mock_updater, name, AsyncMock()) + mock_updater_class.return_value = mock_updater + await executor.execute(request_context, mock_event_queue) + return cast(MagicMock, executor._agent.run) + + @staticmethod + def _forwarded(**kwargs: Any) -> dict[str, Any]: + return {"agent_framework": kwargs} + + async def test_accepted_keys_are_passed_to_agent_run( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor( + mock_agent, accepted_function_invocation_kwargs=["tenant"], accepted_client_kwargs=["trace"] + ) + metadata = self._forwarded( + function_invocation_kwargs={"tenant": "acme", "other": 1}, + client_kwargs={"trace": {"id": "abc"}, "other": 2}, + ) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + kwargs = run.call_args.kwargs + assert kwargs["function_invocation_kwargs"] == {"tenant": "acme"} + assert kwargs["client_kwargs"] == {"trace": {"id": "abc"}} + + async def test_each_destination_is_authorized_independently( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor(mock_agent, accepted_function_invocation_kwargs=["tenant"]) + metadata = self._forwarded( + function_invocation_kwargs={"tenant": "acme"}, + client_kwargs={"tenant": "injected"}, + ) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert run.call_args.kwargs["function_invocation_kwargs"] == {"tenant": "acme"} + assert "client_kwargs" not in run.call_args.kwargs + + async def test_client_kwargs_require_their_own_opt_in( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor(mock_agent, accepted_client_kwargs=["trace"]) + metadata = self._forwarded( + function_invocation_kwargs={"trace": "tool"}, + client_kwargs={"trace": "abc"}, + ) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert "function_invocation_kwargs" not in run.call_args.kwargs + assert run.call_args.kwargs["client_kwargs"] == {"trace": "abc"} + + async def test_nothing_is_accepted_by_default( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor(mock_agent) + metadata = self._forwarded(function_invocation_kwargs={"tenant": "acme"}, client_kwargs={"trace": "abc"}) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert "function_invocation_kwargs" not in run.call_args.kwargs + assert "client_kwargs" not in run.call_args.kwargs + + async def test_configured_run_kwargs_win_on_conflict( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor( + mock_agent, + run_kwargs={"function_invocation_kwargs": {"tenant": "server", "fixed": True}}, + accepted_function_invocation_kwargs=["tenant", "extra"], + ) + metadata = self._forwarded(function_invocation_kwargs={"tenant": "caller", "extra": "x"}) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert run.call_args.kwargs["function_invocation_kwargs"] == {"tenant": "server", "fixed": True, "extra": "x"} + + @mark.parametrize( + "reserved", + ["stream", "messages", "options", "session", "middleware", "client_kwargs", "_internal"], + ) + def test_framework_owned_names_cannot_be_accepted(self, mock_agent: MagicMock, reserved: str) -> None: + with raises(ValueError, match=reserved): + A2AExecutor(mock_agent, accepted_function_invocation_kwargs=["tenant", reserved]) + with raises(ValueError, match=reserved): + A2AExecutor(mock_agent, accepted_client_kwargs=["tenant", reserved]) + + @mark.parametrize( + "metadata", + [{}, {"agent_framework": "text"}, {"agent_framework": {"client_kwargs": "text"}}], + ids=["missing", "not-a-mapping", "values-not-a-mapping"], + ) + async def test_malformed_metadata_is_ignored( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task, metadata: Any + ) -> None: + executor = A2AExecutor( + mock_agent, accepted_function_invocation_kwargs=["tenant"], accepted_client_kwargs=["tenant"] + ) + + run = await self._execute(executor, metadata, mock_event_queue, mock_task) + + assert "client_kwargs" not in run.call_args.kwargs + assert "function_invocation_kwargs" not in run.call_args.kwargs + + async def test_accepted_function_invocation_kwargs_are_passed_when_streaming( + self, mock_agent: MagicMock, mock_event_queue: MagicMock, mock_task: Task + ) -> None: + executor = A2AExecutor(mock_agent, stream=True, accepted_function_invocation_kwargs=["tenant"]) + stream = MagicMock() + stream.with_transform_hook = MagicMock(return_value=stream) + stream.get_final_response = AsyncMock() + request_context = MagicMock() + request_context.get_user_input = MagicMock(return_value="Hello") + request_context.current_task = mock_task + request_context.context_id = "ctx-123" + request_context.message = MagicMock() + request_context.metadata = self._forwarded(function_invocation_kwargs={"tenant": "acme"}) + cast(Any, mock_agent).run = MagicMock(return_value=stream) + cast(Any, mock_agent).create_session = MagicMock() + + with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class: + mock_updater = MagicMock() + for name in ("submit", "start_work", "complete", "update_status"): + setattr(mock_updater, name, AsyncMock()) + mock_updater_class.return_value = mock_updater + await executor.execute(request_context, mock_event_queue) + + assert mock_agent.run.call_args.kwargs["function_invocation_kwargs"] == {"tenant": "acme"}