Python: A2AExecutor: accept forwarded run kwargs from request metadata

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,19 +105,45 @@ 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:
             agent: The AI agent or workflow 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.
                 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):

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"}